nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/rope.py

@@ -0,0 +1,142 @@
+"""RoPE (Rotary Position Embedding) — pure @triton.jit kernels.
+
+Fused forward kernel: applies rotary embedding in-place to Q and K.
+HuggingFace Llama/Mistral style: split-half rotation (not interleaved).
+Extracted from Liger-Kernel (BSD-2-Clause; THIRD_PARTY_NOTICES.md).
+
+Layout:
+    q: (bsz, seq_len, n_q_heads, head_dim)  — physical (transposed from [B,H,S,D])
+    k: (bsz, seq_len, n_kv_heads, head_dim)
+    cos: (1|bsz, seq_len, head_dim)
+    sin: (1|bsz, seq_len, head_dim)
+
+Grid: (bsz * seq_len,)
+Each program handles one (batch, position) across ALL heads.
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def rope_forward_kernel(
+    q_ptr,
+    q_row_stride,
+    k_ptr,
+    k_row_stride,
+    cos,
+    cos_row_stride,
+    sin,
+    sin_row_stride,
+    sl,
+    bs: tl.constexpr,
+    cos_bs: tl.constexpr,
+    n_qh: tl.constexpr,
+    n_kh: tl.constexpr,
+    hd: tl.constexpr,
+    pad_n_qh: tl.constexpr,
+    pad_n_kh: tl.constexpr,
+    pad_hd: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    BACKWARD_PASS: tl.constexpr = False,
+):
+    """Fused RoPE kernel — applies rotary embedding in-place to Q and K.
+
+    Each program instance processes one (batch, seq_position) pair across
+    all Q heads and K heads simultaneously. The rotation formula is:
+
+        y_left  = x_left  * cos - x_right * sin
+        y_right = x_right * cos + x_left  * sin
+
+    For backward pass, sin is negated (equivalent to inverse rotation).
+    """
+    pid = tl.program_id(0).to(tl.int64)
+
+    # Advance pointers to current (batch, position) row
+    q_ptr = q_ptr + pid * q_row_stride
+    k_ptr = k_ptr + pid * k_row_stride
+
+    # Locate cos/sin row for this token position
+    # pid indexes [bsz, seq_len] flattened; seq_len is fastest-changing
+    batch_idx = pid // sl
+    cos_row_idx = pid % sl
+    cos = cos + tl.where(
+        cos_bs == 1,
+        cos_row_idx * cos_row_stride,
+        batch_idx * (sl * cos_row_stride) + cos_row_idx * cos_row_stride,
+    )
+    sin = sin + tl.where(
+        cos_bs == 1,
+        cos_row_idx * sin_row_stride,
+        batch_idx * (sl * sin_row_stride) + cos_row_idx * sin_row_stride,
+    )
+
+    # Load cos/sin (only need left half — right half is identical)
+    cos_offsets = tl.arange(0, pad_hd // 2)
+    cos_mask = cos_offsets < hd // 2
+    # A table WIDER than Q is rounded to Q's dtype on load (round-to-nearest,
+    # the conversion a stored cast makes): the wrapper no longer materialises
+    # a cast copy of the step's fp32 cos/sin per layer for an fp16 Q. A table
+    # narrower than Q is widened by the wrapper beforehand (a copy, said in
+    # its line): widened on load, the fp32 rotation's fused multiply-add
+    # moved with the convert before it (Sana's Gemma-2 encoder, 2026-09-07).
+    cos_row = tl.load(cos + cos_offsets, mask=cos_mask, other=0).to(q_ptr.dtype.element_ty)
+    sin_row = tl.load(sin + cos_offsets, mask=cos_mask, other=0).to(q_ptr.dtype.element_ty)
+
+    # --- Q heads: load left half and right half ---
+    first_half_q_offsets = (
+        tl.arange(0, pad_n_qh)[:, None] * hd
+        + tl.arange(0, pad_hd // 2)[None, :]
+    )
+    first_q_mask = (tl.arange(0, pad_n_qh)[:, None] < n_qh) & (
+        tl.arange(0, pad_hd // 2)[None, :] < hd // 2
+    )
+    q_tile_1 = tl.load(
+        q_ptr + first_half_q_offsets, mask=first_q_mask, other=0
+    ).to(sin_row.dtype)
+
+    second_half_q_offsets = first_half_q_offsets + (hd // 2)
+    q_tile_2 = tl.load(
+        q_ptr + second_half_q_offsets, mask=first_q_mask, other=0
+    ).to(sin_row.dtype)
+
+    # --- K heads: load left half and right half ---
+    first_half_k_offsets = (
+        tl.arange(0, pad_n_kh)[:, None] * hd
+        + tl.arange(0, pad_hd // 2)[None, :]
+    )
+    first_k_mask = (tl.arange(0, pad_n_kh)[:, None] < n_kh) & (
+        tl.arange(0, pad_hd // 2)[None, :] < hd // 2
+    )
+    k_tile_1 = tl.load(
+        k_ptr + first_half_k_offsets, mask=first_k_mask, other=0
+    ).to(sin_row.dtype)
+
+    second_half_k_offsets = first_half_k_offsets + (hd // 2)
+    k_tile_2 = tl.load(
+        k_ptr + second_half_k_offsets, mask=first_k_mask, other=0
+    ).to(sin_row.dtype)
+
+    # --- Apply rotation ---
+    if not BACKWARD_PASS:
+        # Forward: y = [x1, x2] * [cos, cos] + [-x2, x1] * [sin, sin]
+        # The plain expression, as the backend contracts it per dtype: written
+        # with an explicit fma it matched the fp32 path but not the fp16 one
+        # (VibeVoice, 2026-09-07) — the contraction is not the same in both,
+        # and the bytes of every certified row depend on it staying as it is.
+        new_q_1 = q_tile_1 * cos_row - q_tile_2 * sin_row
+        new_q_2 = q_tile_2 * cos_row + q_tile_1 * sin_row
+        new_k_1 = k_tile_1 * cos_row - k_tile_2 * sin_row
+        new_k_2 = k_tile_2 * cos_row + k_tile_1 * sin_row
+    else:
+        # Backward (inverse rotation): negate sin
+        new_q_1 = q_tile_1 * cos_row + q_tile_2 * sin_row
+        new_q_2 = q_tile_2 * cos_row - q_tile_1 * sin_row
+        new_k_1 = k_tile_1 * cos_row + k_tile_2 * sin_row
+        new_k_2 = k_tile_2 * cos_row - k_tile_1 * sin_row
+
+    # Store results in-place
+    tl.store(q_ptr + first_half_q_offsets, new_q_1, mask=first_q_mask)
+    tl.store(q_ptr + second_half_q_offsets, new_q_2, mask=first_q_mask)
+    tl.store(k_ptr + first_half_k_offsets, new_k_1, mask=first_k_mask)
+    tl.store(k_ptr + second_half_k_offsets, new_k_2, mask=first_k_mask)

nbxtensor/ops/rsqrt.py

@@ -0,0 +1,17 @@
+"""Reciprocal square root — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def rsqrt_forward_kernel(
+    input_ptr, output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0).to(tl.int64)
+    offset = pid.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id: the product itself wraps past 2^31 elements (register 58)
+    mask = offset < n_elements
+
+    x = tl.load(input_ptr + offset, mask=mask)
+    tl.store(output_ptr + offset, tl.rsqrt(x.to(tl.float32)), mask=mask)

nbxtensor/ops/scaled_softmax.py

@@ -0,0 +1,96 @@
+"""Scaled softmax — pure @triton.jit kernel.
+
+Computes softmax(input * scale_factor) for attention score tensors.
+Input shape: [B, H, Q, K] (batch, heads, query_seq, key_seq).
+
+Uses online softmax (two-pass): first pass computes max and exp-sum,
+second pass normalizes. Numerically stable.
+
+Extracted from FlagGems reference (scaled_softmax.py).
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def scaled_softmax_kernel(
+    output_ptr,
+    input_ptr,
+    scale_factor,
+    query_seq_len,
+    key_seq_len,
+    stride_b,
+    stride_h,
+    stride_q,
+    BLOCK_Q: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    """Scaled softmax: out = softmax(input * scale_factor) along last dim.
+
+    Grid: (cdiv(query_seq_len, BLOCK_Q), num_heads, batch_size)
+    """
+    query_tile_idx = tl.program_id(0).to(tl.int64)
+    head_idx = tl.program_id(1).to(tl.int64)
+    batch_idx = tl.program_id(2).to(tl.int64)
+
+    start_q = query_tile_idx * BLOCK_Q
+    q_offsets = start_q + tl.arange(0, BLOCK_Q)
+    q_mask = q_offsets < query_seq_len
+
+    row_start = (
+        input_ptr
+        + batch_idx * stride_b
+        + head_idx * stride_h
+        + q_offsets * stride_q
+    )
+
+    # Pass 1: find row max and exp-sum (online softmax)
+    m = tl.full([BLOCK_Q], -float('inf'), dtype=tl.float32)
+    exp_sum = tl.zeros([BLOCK_Q], dtype=tl.float32)
+
+    for k_idx in range(0, tl.cdiv(key_seq_len, BLOCK_K)):
+        k_offsets = k_idx * BLOCK_K + tl.arange(0, BLOCK_K)
+        block_ptr = row_start[:, None] + k_offsets[None, :]
+
+        row_mask = q_mask[:, None]
+        col_mask = k_offsets[None, :] < key_seq_len
+        valid = row_mask & col_mask
+
+        s = tl.load(block_ptr, mask=valid, other=-float('inf'))
+        s = s * scale_factor
+
+        m_new = tl.max(s, axis=1)
+        m_old = m
+        m = tl.maximum(m_old, m_new)
+
+        # Rescale previous sum
+        exp_sum = exp_sum * tl.exp(m_old - m)
+        # Add current block
+        s_exp = tl.exp(s - m[:, None])
+        exp_sum = exp_sum + tl.sum(tl.where(valid, s_exp, 0.0), axis=1)
+
+    exp_sum_inv = 1.0 / exp_sum
+
+    # Pass 2: compute and store normalized values
+    out_row_start = (
+        output_ptr
+        + batch_idx * stride_b
+        + head_idx * stride_h
+        + q_offsets * stride_q
+    )
+
+    for k_idx in range(0, tl.cdiv(key_seq_len, BLOCK_K)):
+        k_offsets = k_idx * BLOCK_K + tl.arange(0, BLOCK_K)
+        in_ptr = row_start[:, None] + k_offsets[None, :]
+        out_ptr = out_row_start[:, None] + k_offsets[None, :]
+
+        row_mask = q_mask[:, None]
+        col_mask = k_offsets[None, :] < key_seq_len
+        valid = row_mask & col_mask
+
+        s = tl.load(in_ptr, mask=valid, other=-float('inf'))
+        s = s * scale_factor
+        p = tl.exp(s - m[:, None]) * exp_sum_inv[:, None]
+
+        tl.store(out_ptr, p, mask=valid)

nbxtensor/ops/scatter_op.py

@@ -0,0 +1,139 @@
+"""Scatter along dim — pure @triton.jit kernel.
+
+Approach from FlagGems: work on flattened index space. For each element,
+decompose into (outer, dim, inner), load the scatter index, then store
+(or atomic_add) into the destination at the indexed position.
+
+scatter: dst[outer][index[outer][dim][inner]][inner] = src[outer][dim][inner]
+scatter_add: dst[...] += src[...]
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def scatter_kernel(
+    src, index, out,
+    N,
+    out_dim_stride,
+    src_stride_outer, src_stride_dim, src_stride_inner,
+    idx_stride_outer, idx_stride_dim, idx_stride_inner,
+    out_stride_outer, out_stride_inner,
+    outer_size, dim_size, inner_size,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Scatter: out[outer][index_val][inner] = src[outer][dim][inner]."""
+    pid = tl.program_id(0).to(tl.int64)
+    offset = pid.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id: the product itself wraps past 2^31 elements (register 58)
+    mask = offset < N
+
+    inner_idx = offset % inner_size
+    temp = offset // inner_size
+    dim_idx = temp % dim_size
+    outer_idx = temp // dim_size
+
+    # Load scatter index
+    idx_offset = outer_idx * idx_stride_outer + dim_idx * idx_stride_dim + inner_idx * idx_stride_inner
+    cur_index = tl.load(index + idx_offset, mask=mask, other=0)
+
+    # Load source value
+    src_offset = outer_idx * src_stride_outer + dim_idx * src_stride_dim + inner_idx * src_stride_inner
+    val = tl.load(src + src_offset, mask=mask, other=0.0)
+
+    # Store to output at scattered position
+    out_offset = outer_idx * out_stride_outer + cur_index * out_dim_stride + inner_idx * out_stride_inner
+    tl.store(out + out_offset, val, mask=mask)
+
+@triton.jit
+def scatter_add_kernel(
+    src, index, out,
+    N,
+    out_dim_stride,
+    src_stride_outer, src_stride_dim, src_stride_inner,
+    idx_stride_outer, idx_stride_dim, idx_stride_inner,
+    out_stride_outer, out_stride_inner,
+    outer_size, dim_size, inner_size,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Scatter add: out[outer][index_val][inner] += src[outer][dim][inner]."""
+    pid = tl.program_id(0).to(tl.int64)
+    offset = pid.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id: the product itself wraps past 2^31 elements (register 58)
+    mask = offset < N
+
+    inner_idx = offset % inner_size
+    temp = offset // inner_size
+    dim_idx = temp % dim_size
+    outer_idx = temp // dim_size
+
+    idx_offset = outer_idx * idx_stride_outer + dim_idx * idx_stride_dim + inner_idx * idx_stride_inner
+    cur_index = tl.load(index + idx_offset, mask=mask, other=0)
+
+    src_offset = outer_idx * src_stride_outer + dim_idx * src_stride_dim + inner_idx * src_stride_inner
+    val = tl.load(src + src_offset, mask=mask, other=0.0)
+
+    out_offset = outer_idx * out_stride_outer + cur_index * out_dim_stride + inner_idx * out_stride_inner
+    tl.atomic_add(out + out_offset, val, mask=mask)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# scatter_reduce kernels — amax, amin using Triton native atomics
+# ═══════════════════════════════════════════════════════════════════════════════
+
+@triton.jit
+def scatter_reduce_amax_kernel(
+    src, index, out,
+    N,
+    out_dim_stride,
+    src_stride_outer, src_stride_dim, src_stride_inner,
+    idx_stride_outer, idx_stride_dim, idx_stride_inner,
+    out_stride_outer, out_stride_inner,
+    outer_size, dim_size, inner_size,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """scatter_reduce with reduce='amax': atomic max."""
+    pid = tl.program_id(0).to(tl.int64)
+    offset = pid.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id: the product itself wraps past 2^31 elements (register 58)
+    mask = offset < N
+
+    inner_idx = offset % inner_size
+    temp = offset // inner_size
+    dim_idx = temp % dim_size
+    outer_idx = temp // dim_size
+
+    idx_offset = outer_idx * idx_stride_outer + dim_idx * idx_stride_dim + inner_idx * idx_stride_inner
+    cur_index = tl.load(index + idx_offset, mask=mask, other=0)
+
+    src_offset = outer_idx * src_stride_outer + dim_idx * src_stride_dim + inner_idx * src_stride_inner
+    val = tl.load(src + src_offset, mask=mask, other=0.0)
+
+    out_offset = outer_idx * out_stride_outer + cur_index * out_dim_stride + inner_idx * out_stride_inner
+    tl.atomic_max(out + out_offset, val, mask=mask)
+
+@triton.jit
+def scatter_reduce_amin_kernel(
+    src, index, out,
+    N,
+    out_dim_stride,
+    src_stride_outer, src_stride_dim, src_stride_inner,
+    idx_stride_outer, idx_stride_dim, idx_stride_inner,
+    out_stride_outer, out_stride_inner,
+    outer_size, dim_size, inner_size,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """scatter_reduce with reduce='amin': atomic min."""
+    pid = tl.program_id(0).to(tl.int64)
+    offset = pid.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id: the product itself wraps past 2^31 elements (register 58)
+    mask = offset < N
+
+    inner_idx = offset % inner_size
+    temp = offset // inner_size
+    dim_idx = temp % dim_size
+    outer_idx = temp // dim_size
+
+    idx_offset = outer_idx * idx_stride_outer + dim_idx * idx_stride_dim + inner_idx * idx_stride_inner
+    cur_index = tl.load(index + idx_offset, mask=mask, other=0)
+
+    src_offset = outer_idx * src_stride_outer + dim_idx * src_stride_dim + inner_idx * src_stride_inner
+    val = tl.load(src + src_offset, mask=mask, other=0.0)
+
+    out_offset = outer_idx * out_stride_outer + cur_index * out_dim_stride + inner_idx * out_stride_inner
+    tl.atomic_min(out + out_offset, val, mask=mask)