nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/grid_sampler.py

@@ -0,0 +1,145 @@
+"""aten::grid_sampler_2d — pure Triton 2D grid sampler.
+
+Clean-room implementation from the PyTorch semantics (ATen GridSampler.h /
+torch._decomp grid sampler decomposition — formulas only, no code vendored):
+
+- Coordinate unnormalization:
+    align_corners=True:  ix = (x + 1)/2 * (W - 1)
+    align_corners=False: ix = ((x + 1)*W - 1)/2
+- Bicubic: cubic-convolution weights with A = -0.75 over a 4x4 tap window;
+  the fractional offset is computed on the RAW unnormalized coordinate
+  (floor before padding), padding is applied PER INTEGER TAP (this staging
+  differs from bilinear and changes border behavior).
+- Bilinear: padding (clamp) is applied to the FLOAT coordinate before
+  flooring; 4 corners accumulated with per-corner in-bounds masks.
+- zeros padding: out-of-bounds taps contribute 0 (masked loads);
+  border padding: tap indices clamp into [0, size-1].
+- All coordinate/weight math and accumulation in float32 (ATen opmath),
+  output cast to the input dtype by the wrapper.
+
+Modes follow the ATen integer encoding: interpolation 0=bilinear,
+2=bicubic (1=nearest unsupported → wrapper raises); padding 0=zeros,
+1=border (2=reflection unsupported → wrapper raises).
+
+Parallelization: one program per (batch, output-location, channel-block);
+coordinates and weights are computed once per location, taps are gathered
+as BLOCK_C-wide strided loads (mirrors the CUDA kernel's channel loop).
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _cc1(x, A: tl.constexpr):
+    # |x| <= 1 branch of the cubic convolution kernel
+    return ((A + 2.0) * x - (A + 3.0)) * x * x + 1.0
+
+
+@triton.jit
+def _cc2(x, A: tl.constexpr):
+    # 1 < |x| < 2 branch of the cubic convolution kernel
+    return ((A * x - 5.0 * A) * x + 8.0 * A) * x - 4.0 * A
+
+
+@triton.jit
+def grid_sampler_2d_kernel(
+    inp_ptr, grid_ptr, out_ptr,
+    C, H, W, OUT_SPATIAL,
+    stride_in, stride_ic, stride_ih, stride_iw,
+    stride_gn, stride_gs, stride_gc,
+    stride_on, stride_oc, stride_os,
+    MODE: tl.constexpr,          # 0 = bilinear, 2 = bicubic
+    PAD_BORDER: tl.constexpr,    # 0 = zeros, 1 = border
+    ALIGN: tl.constexpr,         # 0 / 1
+    BLOCK_C: tl.constexpr,
+):
+    pid_s = tl.program_id(0).to(tl.int64)     # output spatial location (oH*oW flattened)
+    pid_c = tl.program_id(1).to(tl.int64)     # channel block
+    pid_n = tl.program_id(2).to(tl.int64)     # batch
+
+    # A program beyond the output extent reads and writes nothing (masked,
+    # no early exit — unstructured control flow has no lowering on every backend).
+    s_valid = pid_s < OUT_SPATIAL
+
+    c_offs = pid_c * BLOCK_C + tl.arange(0, BLOCK_C)
+    c_mask = (c_offs < C) & s_valid
+
+    # grid holds (x, y) pairs, channels-last
+    gx = tl.load(grid_ptr + pid_n * stride_gn + pid_s * stride_gs
+                 + 0 * stride_gc, mask=s_valid, other=0.0).to(tl.float32)
+    gy = tl.load(grid_ptr + pid_n * stride_gn + pid_s * stride_gs
+                 + 1 * stride_gc, mask=s_valid, other=0.0).to(tl.float32)
+
+    Wf = W.to(tl.float32)
+    Hf = H.to(tl.float32)
+    if ALIGN:
+        ix = (gx + 1.0) * 0.5 * (Wf - 1.0)
+        iy = (gy + 1.0) * 0.5 * (Hf - 1.0)
+    else:
+        ix = ((gx + 1.0) * Wf - 1.0) * 0.5
+        iy = ((gy + 1.0) * Hf - 1.0) * 0.5
+
+    inp_base = inp_ptr + pid_n * stride_in + c_offs * stride_ic
+    acc = tl.zeros([BLOCK_C], dtype=tl.float32)
+
+    if MODE == 0:
+        # ── bilinear: pad the FLOAT coordinate first, then corners ──
+        if PAD_BORDER:
+            ix = tl.minimum(tl.maximum(ix, 0.0), Wf - 1.0)
+            iy = tl.minimum(tl.maximum(iy, 0.0), Hf - 1.0)
+        x0 = tl.floor(ix)
+        y0 = tl.floor(iy)
+        tx = ix - x0
+        ty = iy - y0
+        for j in tl.static_range(2):
+            for i in tl.static_range(2):
+                xi = x0 + i
+                yi = y0 + j
+                wgt = (tl.where(i == 0, 1.0 - tx, tx)
+                       * tl.where(j == 0, 1.0 - ty, ty))
+                inb = (xi >= 0) & (xi < Wf) & (yi >= 0) & (yi < Hf)
+                xs = tl.minimum(tl.maximum(xi, 0.0), Wf - 1.0).to(tl.int32)
+                ys = tl.minimum(tl.maximum(yi, 0.0), Hf - 1.0).to(tl.int32)
+                val = tl.load(inp_base + ys * stride_ih + xs * stride_iw,
+                              mask=c_mask & inb, other=0.0).to(tl.float32)
+                acc += tl.where(inb, wgt, 0.0) * val
+    else:
+        # ── bicubic: floor on the RAW coordinate, pad per integer tap ──
+        A: tl.constexpr = -0.75
+        x0 = tl.floor(ix)
+        y0 = tl.floor(iy)
+        tx = ix - x0
+        ty = iy - y0
+        wx0 = _cc2(tx + 1.0, A)
+        wx1 = _cc1(tx, A)
+        wx2 = _cc1(1.0 - tx, A)
+        wx3 = _cc2(2.0 - tx, A)
+        wy0 = _cc2(ty + 1.0, A)
+        wy1 = _cc1(ty, A)
+        wy2 = _cc1(1.0 - ty, A)
+        wy3 = _cc2(2.0 - ty, A)
+        for j in tl.static_range(4):
+            wy = tl.where(j == 0, wy0,
+                 tl.where(j == 1, wy1,
+                 tl.where(j == 2, wy2, wy3)))
+            yi = y0 + (j - 1)
+            for i in tl.static_range(4):
+                wx = tl.where(i == 0, wx0,
+                     tl.where(i == 1, wx1,
+                     tl.where(i == 2, wx2, wx3)))
+                xi = x0 + (i - 1)
+                inb = (xi >= 0) & (xi < Wf) & (yi >= 0) & (yi < Hf)
+                xs = tl.minimum(tl.maximum(xi, 0.0), Wf - 1.0).to(tl.int32)
+                ys = tl.minimum(tl.maximum(yi, 0.0), Hf - 1.0).to(tl.int32)
+                if PAD_BORDER:
+                    val = tl.load(inp_base + ys * stride_ih + xs * stride_iw,
+                                  mask=c_mask, other=0.0).to(tl.float32)
+                    acc += wy * wx * val
+                else:
+                    val = tl.load(inp_base + ys * stride_ih + xs * stride_iw,
+                                  mask=c_mask & inb, other=0.0).to(tl.float32)
+                    acc += tl.where(inb, wy * wx, 0.0) * val
+
+    out = out_ptr + pid_n * stride_on + c_offs * stride_oc + pid_s * stride_os
+    tl.store(out, acc, mask=c_mask)

nbxtensor/ops/group_norm_tiled.py

@@ -0,0 +1,67 @@
+"""The two passes of an exact tiled group norm on the Triton engine (`core/strategies/tiled_norm.py`'s NBXTensor twin):
+the statistics of each (sample, group) over the WHOLE input, then a slice normalised with them.
+
+`group_norm_stats_kernel` is the first pass of `groupnorm.group_norm_forward_kernel` alone — the same fp32 sum and
+sum of squares, var = ss/n - mean^2 — without writing the normalised tensor (a stretch sliced along the axis holds its
+input whole, never a whole output beside it). `group_norm_apply_kernel` is its second pass on a slice, reading the
+given statistics. One program per (sample, group); the walk is a fixed chunk (no autotune key: nothing to certify), so
+the summation order is fixed.
+
+Pure @triton.jit (R33).
+"""
+
+import triton
+import triton.language as tl
+
+#: the chunk a program walks its group by — fixed, so the statistics' summation order is
+TILED_NORM_BLOCK = 4096
+TILED_NORM_WARPS = 8
+
+
+@triton.jit
+def group_norm_stats_kernel(X, Mean, Rstd, group_size, C, HW, num_groups, eps, BLOCK_SIZE: tl.constexpr):
+    pid = tl.program_id(0).to(tl.int64)
+    batch_idx = pid // num_groups
+    group_idx = pid % num_groups
+    base = batch_idx * C * HW + group_idx * group_size * HW
+    hidden = group_size * HW
+    block = tl.arange(0, BLOCK_SIZE)
+    s = tl.zeros((), dtype=tl.float32)
+    ss = tl.zeros((), dtype=tl.float32)
+    for off in tl.range(0, hidden, BLOCK_SIZE):
+        idx = off + block
+        mask = idx < hidden
+        x = tl.load(X + base + idx, mask=mask, other=0.0).to(tl.float32)
+        s += tl.sum(x)
+        ss += tl.sum(x * x)
+    inv_n = 1.0 / hidden
+    mean = s * inv_n
+    var = ss * inv_n - mean * mean
+    tl.store(Mean + pid, mean)
+    tl.store(Rstd + pid, 1.0 / tl.sqrt(var + eps))
+
+
+@triton.jit
+def group_norm_apply_kernel(X, Y, W, B, Mean, Rstd, group_size, C, HW, num_groups,
+                            scale_by_weight: tl.constexpr, add_bias: tl.constexpr, BLOCK_SIZE: tl.constexpr):
+    pid = tl.program_id(0).to(tl.int64)
+    batch_idx = pid // num_groups
+    group_idx = pid % num_groups
+    chan_start = group_idx * group_size
+    base = batch_idx * C * HW + chan_start * HW
+    hidden = group_size * HW
+    mean = tl.load(Mean + pid)
+    rstd = tl.load(Rstd + pid)
+    block = tl.arange(0, BLOCK_SIZE)
+    for off in tl.range(0, hidden, BLOCK_SIZE):
+        idx = off + block
+        mask = idx < hidden
+        x = tl.load(X + base + idx, mask=mask, other=0.0).to(tl.float32)
+        y = (x - mean) * rstd
+        if scale_by_weight or add_bias:
+            chan = chan_start + idx // HW
+            if scale_by_weight:
+                y = y * tl.load(W + chan, mask=mask, other=1.0).to(tl.float32)
+            if add_bias:
+                y = y + tl.load(B + chan, mask=mask, other=0.0).to(tl.float32)
+        tl.store(Y + base + idx, y.to(Y.dtype.element_ty), mask=mask)

nbxtensor/ops/groupnorm.py

@@ -0,0 +1,123 @@
+"""Group normalization — pure @triton.jit kernel.
+
+Two-pass loop over hidden_size (= group_size * HW) in BLOCK_SIZE chunks,
+so the kernel scales to arbitrary spatial dimensions. Replaces the
+original FlagGems extraction whose single-tile design crashed with
+`numel exceeds triton maximum tensor numel (1048576)` once HW reached
+1024×1024 (PixArt VAE 1024px output).
+
+Algorithm: standard sum-of-squares variance per (batch, group) program.
+For each program (one (batch, group) pair):
+  1. Loop over hidden_size in BLOCK_SIZE chunks → accumulate sum and
+     sum-of-squares → mean = s / N, var = ss/N - mean².
+  2. rstd = 1 / sqrt(var + eps).
+  3. Loop a second time → normalize → multiply by per-channel weight,
+     add per-channel bias, store.
+
+Forward only (inference). The wrapper guarantees BLOCK_SIZE ≤
+triton.next_power_of_2(hidden_size_per_chunk_target) and ≤ 16384, well
+under Triton's 2^20 numel ceiling on every architecture.
+"""
+
+import triton
+import triton.language as tl
+from triton import next_power_of_2
+
+from nbxtensor.ops._configs import nbx_autotune
+
+
+def _group_norm_configs():
+    """The candidates the certifier proves (chantier 1): the chunk a group is walked by x warps. The one the kernel ran
+    before certification (the group's padded extent capped at 16 384, 4 warps) is among them."""
+    return [triton.Config({'BLOCK_SIZE': block}, num_warps=warps)
+            for block in (1024, 2048, 4096, 8192, 16384) for warps in (2, 4, 8, 16)]
+
+
+def _prune_group_norm(configs, named_args, **kwargs):
+    """A chunk no wider than the backend's reduce-tile cap (`nbx_tensor.reduce_tile_max`) nor than the group's padded
+    extent (a wider chunk only masks)."""
+    from nbxtensor.backend.capabilities import reduce_tile_max
+    hidden = int(named_args['group_size']) * int(named_args['HW'])
+    cap = min(reduce_tile_max(), next_power_of_2(max(1, hidden)))
+    keep = [c for c in configs if c.kwargs['BLOCK_SIZE'] <= cap]
+    return keep or sorted(configs, key=lambda c: c.kwargs['BLOCK_SIZE'])[:1]
+
+
+@nbx_autotune(configs=_group_norm_configs(), key=['PROG_BUCKET', 'HIDDEN_BUCKET', 'scale_by_weight', 'add_bias'],
+              prune_configs_by={'early_config_prune': _prune_group_norm}, cache_results=True)
+@triton.jit
+def group_norm_forward_kernel(
+    X,                              # input pointer, shape (N, C, HW)
+    Y,                              # output pointer, shape (N, C, HW)
+    W,                              # weight pointer, shape (C,)
+    B,                              # bias pointer, shape (C,)
+    Mean,                           # mean output, shape (N * num_groups,)
+    Rstd,                           # rstd output, shape (N * num_groups,)
+    group_size,                     # channels per group
+    C,                              # total channels
+    HW,                             # spatial size = H * W (or H * W * D)
+    num_groups,                     # total groups
+    eps,
+    PROG_BUCKET,                    # the programs' (batch x groups) bucket: the key
+    HIDDEN_BUCKET,                  # the group's padded extent: the key
+    scale_by_weight: tl.constexpr,
+    add_bias: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,       # chunk size along hidden_size: the certified configuration
+):
+    """One program handles one (batch, group) pair. Spatial dimension is
+    walked in BLOCK_SIZE chunks so the per-program tile stays bounded
+    regardless of input HW."""
+    pid = tl.program_id(0).to(tl.int64)
+    batch_idx = pid // num_groups
+    group_idx = pid % num_groups
+
+    # Channel range of this group (each channel has HW elements contiguous
+    # in memory; the group's elements are group_size*HW elements starting
+    # at base_offset).
+    chan_start = group_idx * group_size
+    base_offset = batch_idx * C * HW + chan_start * HW
+    hidden = group_size * HW
+
+    block_range = tl.arange(0, BLOCK_SIZE)
+
+    # ── Pass 1: sum and sum-of-squares over the full hidden dim ─────────
+    s = tl.zeros((), dtype=tl.float32)
+    ss = tl.zeros((), dtype=tl.float32)
+    for off in tl.range(0, hidden, BLOCK_SIZE):
+        idx = off + block_range
+        mask = idx < hidden
+        x = tl.load(X + base_offset + idx, mask=mask, other=0.0).to(tl.float32)
+        s += tl.sum(x)
+        ss += tl.sum(x * x)
+
+    inv_n = 1.0 / hidden
+    mean = s * inv_n
+    var = ss * inv_n - mean * mean
+    rstd = 1.0 / tl.sqrt(var + eps)
+
+    tl.store(Mean + pid, mean)
+    tl.store(Rstd + pid, rstd)
+
+    # ── Pass 2: normalize → optional affine → store ──────────────────────
+    # Each chunk falls inside a single channel because chunks walk the
+    # contiguous (channel, HW) layout. We compute the channel for each
+    # element of the block to fetch the right (W, B) entries.
+    for off in tl.range(0, hidden, BLOCK_SIZE):
+        idx = off + block_range
+        mask = idx < hidden
+        x = tl.load(X + base_offset + idx, mask=mask, other=0.0).to(tl.float32)
+        x_hat = (x - mean) * rstd
+
+        if scale_by_weight or add_bias:
+            # Channel index within the GROUP, then add the group's chan_start.
+            chan_in_group = idx // HW
+            chan_global = chan_start + chan_in_group
+            chan_mask = mask & (chan_in_group < group_size)
+            if scale_by_weight:
+                w = tl.load(W + chan_global, mask=chan_mask, other=1.0).to(tl.float32)
+                x_hat = x_hat * w
+            if add_bias:
+                b = tl.load(B + chan_global, mask=chan_mask, other=0.0).to(tl.float32)
+                x_hat = x_hat + b
+
+        tl.store(Y + base_offset + idx, x_hat, mask=mask)

nbxtensor/ops/gt.py

@@ -0,0 +1,32 @@
+"""Greater than — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def gt_forward_kernel(
+    x_ptr, y_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(x_ptr + offset, mask=mask)
+    y = tl.load(y_ptr + offset, mask=mask)
+    tl.store(output_ptr + offset, (x > y).to(tl.int1), mask=mask)
+
+@triton.jit
+def gt_scalar_kernel(
+    x_ptr, output_ptr,
+    n_elements,
+    scalar,
+    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(x_ptr + offset, mask=mask)
+    tl.store(output_ptr + offset, (x > scalar).to(tl.int1), mask=mask)

nbxtensor/ops/hardsigmoid.py

@@ -0,0 +1,18 @@
+"""Hard Sigmoid — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def hardsigmoid_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)
+    out = tl.maximum(0, tl.minimum(1, x / 6.0 + 0.5))
+    tl.store(output_ptr + offset, out, mask=mask)

nbxtensor/ops/hardswish.py

@@ -0,0 +1,19 @@
+"""Hard Swish — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def hardswish_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)
+    # hardswish = x * relu6(x + 3) / 6
+    out = x * tl.minimum(tl.maximum(0, x + 3.0), 6.0) / 6.0
+    tl.store(output_ptr + offset, out, mask=mask)

nbxtensor/ops/im2col.py

@@ -0,0 +1,57 @@
+"""aten::im2col (nn.Unfold) — pure @triton.jit kernel.
+
+Extracts sliding local blocks from a batched NCHW tensor into columns:
+    input  [N, C, IH, IW]
+    output [N, C*KH*KW, L]   with L = OH*OW
+
+    out[n, c*KH*KW + kh*KW + kw, oh*OW + ow]
+        = input[n, c, oh*stride_h - pad_h + kh*dil_h,
+                       ow*stride_w - pad_w + kw*dil_w]   (0 outside the image)
+
+This is the operation behind HAT's OCAB overlapping-window cross-attention
+(P-TRITON-IM2COL-KERNEL). Pure Triton (R33): NBXTensor in/out, tl.load with
+boundary masking for the padding halo (no F.pad).
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def im2col_kernel(
+    inp_ptr, out_ptr,
+    C, IH, IW,
+    OW, L,
+    KH: tl.constexpr, KW: tl.constexpr,
+    stride_h: tl.constexpr, stride_w: tl.constexpr,
+    pad_h: tl.constexpr, pad_w: tl.constexpr,
+    dil_h: tl.constexpr, dil_w: tl.constexpr,
+    in_sn, in_sc, in_sh, in_sw,
+    out_sn, out_sc, out_sl,
+    BLOCK_L: tl.constexpr,
+):
+    # One program per (n, output-channel) = (n, c*KH*KW + kh*KW + kw); it
+    # writes one full output row of L block-positions (looped in BLOCK_L tiles).
+    pid = tl.program_id(0).to(tl.int64)
+    CKK = C * KH * KW
+    n = pid // CKK
+    oc = pid % CKK
+    c = oc // (KH * KW)
+    rem = oc % (KH * KW)
+    kh = rem // KW
+    kw = rem % KW
+
+    in_base = n * in_sn + c * in_sc
+    out_base = n * out_sn + oc * out_sc
+
+    for l0 in range(0, L, BLOCK_L):
+        l = l0 + tl.arange(0, BLOCK_L)
+        mask_l = l < L
+        oh = l // OW
+        ow = l % OW
+        ih = oh * stride_h - pad_h + kh * dil_h
+        iw = ow * stride_w - pad_w + kw * dil_w
+        in_bounds = mask_l & (ih >= 0) & (ih < IH) & (iw >= 0) & (iw < IW)
+        vals = tl.load(inp_ptr + in_base + ih * in_sh + iw * in_sw,
+                       mask=in_bounds, other=0.0)
+        tl.store(out_ptr + out_base + l * out_sl, vals, mask=mask_l)

nbxtensor/ops/index_add.py

@@ -0,0 +1,95 @@
+"""Index add along dim — deterministic, pure @triton.jit.
+
+out[..., index[i], ...] += alpha * src[..., i, ...]
+
+The previous implementation used `tl.atomic_add`, whose inter-block
+float accumulation order is non-deterministic. When several source
+positions map to the same destination (MoE expert aggregation: many
+routed rows -> one token), repeated runs produced different fp32
+roundings -> a different greedy argmax on Qwen3-30B::triton's unfused
+MoE path (workstream P-TRITON-MOE-DETERMINISM sub-workstream 2).
+
+State of the art consulted (R16, 2026-05-19):
+  - Triton #7402: tl.atomic_add return value is incorrect across
+    threads (only thread 0 correct) — open upstream bug.
+  - Triton #4717: tl.atomic_add is slow (layout conversions).
+  - flagos-ai/FlagGems ops/index_add.py still uses
+    tl.atomic_add(sem='relaxed') (pins triton==3.3.0, no bit-exact
+    test) — no deterministic upstream to vendor.
+  - PyTorch index_add_/scatter_add_ are non-deterministic on CUDA via
+    atomicAdd; the deterministic path is sort + sequential segmented
+    reduce.
+
+Strategy — output-owner gather (deterministic by construction):
+invert the scatter. Each output cell is owned by exactly ONE program
+that sums its contributing source rows in a fixed ascending source-
+position order as a SEQUENTIAL fp32 fold. No atomics, no inter-block
+races, no sort dependency. The ascending-j sequential fold is exactly
+the order torch's deterministic CUDA index_add_ uses (verified
+bit-exact fp32 for bucket sizes up to 4096). A pairwise/tree reduce
+(`tl.sum`) would round differently and break bit-exactness, so the
+j accumulation is a scalar loop; the feature (inner) dim is vectorised.
+
+Out-of-range source indices never equal a valid destination d in
+[0, inp_shape_dim) so they are naturally dropped (mirrors the old
+masked-store semantics — no assert/crash). Model-agnostic, no autotune.
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def index_add_gather_kernel(
+    x, src, index, out,
+    dim_size, inp_shape_dim, inner_size, alpha,
+    src_stride_outer, src_stride_dim, src_stride_inner,
+    INNER_BLOCK: tl.constexpr,
+):
+    """One program == one (outer, dest, inner-tile) output cell.
+
+    out[o, d, i] = x[o, d, i]
+                 + alpha * sum_{j : index[j] == d} src[o, j, i]
+    with the sum taken in ascending-j order as a sequential fp32 fold
+    (bit-exact vs torch deterministic index_add_).
+    """
+    n_inner_tiles = (inner_size + INNER_BLOCK - 1) // INNER_BLOCK
+    pid = tl.program_id(0).to(tl.int64)
+    tile = pid % n_inner_tiles
+    tmp = pid // n_inner_tiles
+    d = tmp % inp_shape_dim
+    o = tmp // inp_shape_dim
+
+    inner = tile * INNER_BLOCK + tl.arange(0, INNER_BLOCK)
+    inner_mask = inner < inner_size
+
+    out_base = o * (inp_shape_dim * inner_size) + d * inner_size + inner
+    base = tl.load(x + out_base, mask=inner_mask, other=0.0).to(tl.float32)
+
+    # torch's deterministic CUDA index_add_ computes the per-destination
+    # segment sum SEPARATELY (0-initialised, ascending-j sequential
+    # fold) and adds it to `self` ONCE at the end — NOT folding into
+    # self. Verified by discriminating probe (base=1, src=[1e8,-1e8]:
+    # torch -> 1.0 == base + Σ, not 0.0 == fold-into-base).
+    seg = tl.zeros([INNER_BLOCK], dtype=tl.float32)
+    src_o = o * src_stride_outer + inner * src_stride_inner
+    # Scan source positions in ascending j (deterministic order).
+    # Scalar loop on purpose: a vectorised tree reduce would round
+    # differently and break bit-exactness vs torch.
+    for j in range(0, dim_size):
+        dest = tl.load(index + j).to(tl.int64)
+        if dest == d:
+            sv = tl.load(src + src_o + j * src_stride_dim,
+                         mask=inner_mask, other=0.0).to(tl.float32)
+            seg += (alpha * sv).to(tl.float32)
+            # NOTE alpha != 1: a residual 1-ULP fp32 gap vs torch
+            # remains (Triton codegen of the scaled term — bitcast /
+            # fma / reassoc barriers were tried, all no-effect). It is
+            # NOT an ordering bug: numpy fp32 sequential fold == torch
+            # bit-exact, and this kernel reproduces that for alpha == 1
+            # bit-exact. alpha != 1 never occurs on a transformer path
+            # (Qwen3-30B-A3B graph: 6144/6144 index_add use alpha=1).
+
+    # Triton implicitly casts the fp32 result to out's element type on
+    # store (same convention as the other elementwise kernels).
+    tl.store(out + out_base, base + seg, mask=inner_mask)

nbxtensor/ops/index_put_op.py

@@ -0,0 +1,92 @@
+"""index_put — pure @triton.jit kernel.
+
+Scatter-write `values` into `out` at positions given by a single 1-D
+advanced index tensor on the leading dim (the form decomposed-ATen
+emits for MoE-v2 expert-output aggregation, KV-cache indexed writes,
+and post-`aten::nonzero` masked scatter):
+
+    out[idx[s], t] = values[s, t]              (ACCUMULATE == 0)
+    out[idx[s], t] += values[s, t]             (ACCUMULATE == 1)
+
+`idx` may have any shape S (flattened, numel = Sn); `t` ranges over
+the flattened tail `x.shape[1:]` (numel = T).  `out` is contiguous so
+the leading-dim stride is exactly T.  Duplicate indices follow torch
+semantics: summed when accumulating, last-writer-wins (unspecified
+order) otherwise.
+
+Index decomposition logic mirrors `index_add_op.py` (FlagGems
+lineage, Apache-2.0).
+"""
+
+import triton
+import triton.language as tl
+
+#: The contract this kernel enforces. Triton forbids reading a module global
+#: from inside a @jit body, so the kernel repeats the literal; the permanent
+#: test pins the two equal so they cannot drift.
+INDEX_PUT_OOB = "index_put: index out of range"
+
+
+@triton.jit(debug=True)
+def index_put_kernel(
+    out_ptr, idx_ptr, val_ptr,
+    T, N, R, fault_ptr,
+    FAULT_CODE: tl.constexpr,
+    VAL_SCALAR: tl.constexpr,
+    ACCUMULATE: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """out[idx[off // T], off % T] {=|+=} values[off]  for off < N.
+
+    `R` = rows of `out` (its leading dim). A negative index counts from
+    the end (torch semantics); an index still out of [0, R) after that
+    TRAPS via `tl.device_assert` (kept in the binary by `debug=True`)
+    where torch's index_put raises — the FlagGems-lineage kernel wrote
+    through it to memory outside the tensor, silently
+    (D-GATHER-SCATTER-OOB-SILENT, promoted to correctness 2026-09-02).
+
+    That assert is honoured on CUDA and ROCm and NOT on Metal, which
+    computes its predicate and discards it (measured 2026-09-10). Here
+    that was not a wrong value, it was the WRITE ITSELF: the store's mask
+    was `off < N` alone, so `dst = row * T + t` with an out-of-range row
+    wrote outside the tensor — 8 floats past the end of a 24-float
+    tensor in the eight-float reproducer. Two things close it. The
+    address is made legal before it is used, so nothing outside the
+    tensor is touched in the interval before the refusal arrives (on
+    CUDA the assert halts the thread first, so it changes nothing
+    there); and `FAULT_CODE`, a compile-time constant, carries the
+    refusal itself through the fault word, which `check_device_faults`
+    raises on at the next host observation. The wrapper passes 0 where
+    the assert is honoured, and then not one instruction of it is
+    emitted.
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    off = 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 = off < N
+
+    s = off // T                       # which index entry (row in idx)
+    t = off % T                        # offset within the flattened tail
+    row = tl.load(idx_ptr + s, mask=mask, other=0).to(tl.int64)
+    row = tl.where(row < 0, row + R, row)
+    # Masked lanes loaded `other=0` and are in range by construction.
+    row_valid = (row >= 0) & (row < R)
+    tl.device_assert(row_valid, "index_put: index out of range")
+    if FAULT_CODE != 0:
+        # Lanes past N loaded `other=0`; they must not report (and `R == 0`
+        # would make even row 0 invalid, so `mask` has to gate this).
+        if tl.max(tl.where(mask & (row_valid == 0), 1, 0)) != 0:
+            tl.store(fault_ptr, FAULT_CODE)
+    # A legal address, so no memory outside the tensor is touched before the
+    # refusal is observed. Not a rescue: the run refuses either way.
+    dst = tl.where(row_valid, row, 0) * T + t
+
+    if VAL_SCALAR:
+        # All lanes read element 0 (off * 0 keeps it a vector for masking).
+        v = tl.load(val_ptr + off * 0, mask=mask, other=0.0)
+    else:
+        v = tl.load(val_ptr + off, mask=mask, other=0.0)
+
+    if ACCUMULATE:
+        tl.atomic_add(out_ptr + dst, v, mask=mask & row_valid)
+    else:
+        tl.store(out_ptr + dst, v, mask=mask & row_valid)