nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/paged_decode_attn.py

@@ -0,0 +1,215 @@
+"""Paged KV: the append and the decode attention read through a page table (state engine step 2,
+docs/design/state-engine.md §4.3; the shared bookkeeping is core/state/pages.py).
+
+The state of a decoding sequence lives in fixed-size pages of a per-node pool, granted as the sequence grows and never
+at max_len. One pool per layer and per K / V, laid out [num_pages, kv_heads, page_tokens, D]: the positions of one KV
+head inside one page are contiguous rows, so a tile of BLOCK_N positions of one head is one strided block read.
+
+- `paged_kv_append_kernel` writes new positions through the slot mapping (vLLM's `reshape_and_cache` shape): slot =
+  page * page_tokens + position-in-page, one int32 per new token, computed by the PageTable.
+- `paged_decode_attn_split_kernel` is the vector (SIMT) split-KV decode read of ops/decode_attn_vec.py — one query row
+  per program, FMA reduction over D, scalar online softmax, fp32 accumulator, partials merged by the fixed-order
+  `flash_decode_reduce_kernel` — with K / V addressed through the block table (one page id per page of the sequence,
+  int32, padded with -1). A tile never straddles two pages: every candidate tile divides the page size
+  (`PAGED_DECODE_TILE_CANDIDATES`, ops/tiling_sizes, read by core/state/price.page_tokens) and each split starts on a tile boundary.
+  Each sequence of a batch reads its own length (`seq_lens`), so one launch serves a batch of different lengths.
+
+Pure @triton.jit, no tl.dot (sm_70: the vector kernel's measured structure), one source for every backend.
+"""
+
+import triton
+import triton.language as tl
+
+from nbxtensor.ops.decode_attn_vec import _round_q
+
+
+# The page tiles are sizes: nbxtensor.ops.tiling_sizes owns them (torch- and triton-free — the compiled branch sizes
+# its paged cache with them in an environment without triton).
+from nbxtensor.ops.tiling_sizes import PAGED_DECODE_TILE_CANDIDATES  # noqa: F401
+from nbxtensor.ops._configs import nbx_autotune
+
+# The split shape: positions per split and the most splits one launch makes — the vector decode kernel's judged config
+# (its locked protocol of 2026-08-23, the defaults of `_decode_attn_vec`'s NBX_DV_SEG and its 64-split cap), so the
+# paged read partitions a sequence exactly as the contiguous read it replaces.
+PAGED_DECODE_SPLIT_POSITIONS = 256
+PAGED_DECODE_MAX_SPLITS = 64
+
+
+def paged_decode_splits(longest: int) -> int:
+    """The splits one decode launch makes for its longest sequence — the read's grid, the autotune key's length class
+    (what a split's work depends on) and the replayed decode's bucket (`PagedTritonKVCache.replay_key`): one rule, so
+    a replayed step never meets a launch the certificate keys differently. Keyed on the length's own bucket instead,
+    which moves at every token below 64, a replay recorded at one step was verified against another step's
+    configuration, and the bytes differed: TinyLlama's decode stopped replaying (25 -> 135 ms a token, 2026-11-04)."""
+    return max(1, min(-(-max(int(longest), 1) // PAGED_DECODE_SPLIT_POSITIONS), PAGED_DECODE_MAX_SPLITS))
+
+
+@triton.jit
+def paged_kv_append_kernel(
+    new_ptr,         # [N_tok, H_kv, D] the new positions' K (or V), any float dtype
+    pool_ptr,        # [num_pages, H_kv, page_tokens, D] the layer's pool
+    slot_ptr,        # [N_tok] int32: page * page_tokens + position in page
+    page_tokens,
+    stride_nt, stride_nh,
+    stride_pp, stride_ph, stride_ps,
+    BLOCK_D: tl.constexpr,
+    D: tl.constexpr,
+):
+    """One program per (new token, KV head): its D values to the page slot the slot mapping names, cast to the pool's
+    dtype on store."""
+    pid_t = tl.program_id(0).to(tl.int64)
+    pid_h = tl.program_id(1).to(tl.int64)
+    offs_d = tl.arange(0, BLOCK_D)
+    mask_d = offs_d < D
+    slot = tl.load(slot_ptr + pid_t).to(tl.int64)
+    page = slot // page_tokens
+    pos = slot % page_tokens
+    x = tl.load(new_ptr + pid_t * stride_nt + pid_h * stride_nh + offs_d, mask=mask_d, other=0.0)
+    tl.store(pool_ptr + page * stride_pp + pid_h * stride_ph + pos * stride_ps + offs_d,
+             x.to(pool_ptr.dtype.element_ty), mask=mask_d)
+
+
+#: the position tiles the paged read is certified over (each dividing the page, `_tiles_in_the_page`), and the warps
+_PAGED_DECODE_CONFIGS = [triton.Config({'BLOCK_N': bn}, num_warps=w, num_stages=1)
+                         for bn in PAGED_DECODE_TILE_CANDIDATES for w in (2, 4, 8)]
+#: a split's positions are a whole number of the LARGEST tile, so every candidate tile cuts the splits the same way:
+#: the partials, and the oracle a certification proves them against, are the split's, not the tile's
+PAGED_DECODE_SPLIT_ALIGN = max(PAGED_DECODE_TILE_CANDIDATES)
+
+
+def _tiles_in_the_page(configs, named_args, **kwargs):
+    """The candidates whose tile divides the page (a tile never straddles two pages)."""
+    pt = int(named_args["page_tokens"])
+    return [c for c in configs if pt % int(c.kwargs["BLOCK_N"]) == 0]
+
+
+@nbx_autotune(configs=_PAGED_DECODE_CONFIGS, key=['SPLITS', 'page_tokens', 'D', 'D_V', 'GQA_GROUPS'],
+              prune_configs_by={'early_config_prune': _tiles_in_the_page}, cache_results=True)
+@triton.jit
+def paged_decode_attn_split_kernel(
+    q_ptr,           # [B*H_q, D] one query row per program
+    k_pool_ptr,      # [num_pages, H_kv, page_tokens, D]
+    v_pool_ptr,      # [num_pages, H_kv, page_tokens, D_v]
+    block_table_ptr, # [B, max_pages] int32 page ids, -1 past the sequence
+    seq_lens_ptr,    # [B] int32 positions each sequence holds
+    opart_ptr,       # [B*H_q, SPLIT, D_v] fp32 partial outputs
+    mpart_ptr,       # [B*H_q, SPLIT] fp32 running max per split
+    lpart_ptr,       # [B*H_q, SPLIT] fp32 sum-of-exp per split
+    n_seqs,          # B: the split's length is read from the device's lengths, never passed (see below)
+    page_tokens,     # positions per page, a multiple of BLOCK_N
+    SPLITS,          # the grid's splits (`paged_decode_splits`): the autotune key's length class (unused in the body)
+    sm_scale,
+    H_Q,
+    stride_qh,
+    stride_kp, stride_kh, stride_ks,
+    stride_vp, stride_vh, stride_vs,
+    stride_bt,
+    stride_oh, stride_os,
+    stride_mh,
+    GQA_GROUPS: tl.constexpr,   # H_q // H_kv
+    BLOCK_D: tl.constexpr,
+    D: tl.constexpr,
+    D_V: tl.constexpr,
+    SPLIT_ALIGN: tl.constexpr,        # a split's positions: a whole number of these (PAGED_DECODE_SPLIT_ALIGN)
+    Q_TO: tl.constexpr = None,        # a dtype: Q is rounded to it on load (the cache's), as the contiguous read does
+    Q_SATURATE: tl.constexpr = False,
+    BLOCK_N: tl.constexpr = 32,
+):
+    pid_q = tl.program_id(0).to(tl.int64)        # b * H_q + h_q
+    pid_s = tl.program_id(1).to(tl.int64)        # split index
+    b = pid_q // H_Q
+    h_kv = (pid_q % H_Q) // GQA_GROUPS
+
+    offs_d = tl.arange(0, BLOCK_D)
+    mask_d = offs_d < D
+    mask_dv = offs_d < D_V
+    q = _round_q(tl.load(q_ptr + pid_q * stride_qh + offs_d, mask=mask_d, other=0.0), Q_TO, Q_SATURATE)
+
+    T_b = tl.load(seq_lens_ptr + b).to(tl.int32)
+    # Positions per split, read from the DEVICE lengths: the longest sequence over the grid's splits, in whole tiles
+    # — the host's figure from the same lengths, so the same value at every step. Passed as a scalar it was frozen into
+    # a replayed decode's launch: the split count only moves every PAGED_DECODE_SPLIT_POSITIONS positions, the length
+    # every token, and a replayed TinyLlama stopped attending to its newest positions (2026-10-28).
+    longest = tl.load(seq_lens_ptr).to(tl.int32)
+    for i in range(1, n_seqs):
+        longest = tl.maximum(longest, tl.load(seq_lens_ptr + i).to(tl.int32))
+    longest = tl.maximum(longest, 1)
+    n_split = tl.num_programs(1)
+    seg_len = tl.cdiv(tl.cdiv(longest, n_split), SPLIT_ALIGN) * SPLIT_ALIGN
+    seg_start = (pid_s * seg_len).to(tl.int32)   # 32-bit loop bounds (the Metal lowering refuses a 64-bit scf.for bound)
+    seg_end = tl.minimum(seg_start + seg_len, T_b).to(tl.int32)
+
+    m_i = float("-inf")
+    l_i = 0.0
+    acc = tl.zeros((BLOCK_D,), dtype=tl.float32)
+
+    for n0 in range(seg_start, seg_end, BLOCK_N):
+        # the tile's page: n0 is a multiple of BLOCK_N and BLOCK_N divides page_tokens, so the tile is in one page
+        page = tl.load(block_table_ptr + b * stride_bt + n0 // page_tokens).to(tl.int64)
+        in_page = (n0 % page_tokens) + tl.arange(0, BLOCK_N)
+        mask_n = (n0 + tl.arange(0, BLOCK_N)) < seg_end
+
+        k_tile = tl.load(k_pool_ptr + page * stride_kp + h_kv * stride_kh
+                         + in_page[:, None] * stride_ks + offs_d[None, :],
+                         mask=mask_n[:, None] & mask_d[None, :], other=0.0)
+        s = tl.sum(q[None, :] * k_tile.to(tl.float32), 1) * sm_scale
+        s = tl.where(mask_n, s, float("-inf"))
+
+        m_new = tl.maximum(m_i, tl.max(s, 0))
+        m_safe = tl.where(m_new == float("-inf"), 0.0, m_new)
+        corr = tl.exp(tl.where(m_i == float("-inf"), float("-inf"), m_i - m_safe))
+        p = tl.exp(s - m_safe)
+        p = tl.where(mask_n, p, 0.0)
+
+        v_tile = tl.load(v_pool_ptr + page * stride_vp + h_kv * stride_vh
+                         + in_page[:, None] * stride_vs + offs_d[None, :],
+                         mask=mask_n[:, None] & mask_dv[None, :], other=0.0)
+        acc = acc * corr + tl.sum(p[:, None] * v_tile.to(tl.float32), 0)
+        l_i = l_i * corr + tl.sum(p, 0)
+        m_i = m_new
+
+    tl.store(opart_ptr + pid_q * stride_oh + pid_s * stride_os + offs_d, acc, mask=mask_dv)
+    tl.store(mpart_ptr + pid_q * stride_mh + pid_s, m_i)
+    tl.store(lpart_ptr + pid_q * stride_mh + pid_s, l_i)
+
+
+@triton.jit
+def paged_kv_append_step_kernel(
+    k_new_ptr, v_new_ptr,   # [B, H_kv, D] / [B, H_kv, D_v] the step's one new position per sequence
+    k_pool_ptr, v_pool_ptr, # [num_pages, H_kv, page_tokens, D(_v)]
+    block_table_ptr,        # [B, max_pages] int32
+    lens_ptr,               # [B] int32: the layer's positions held — the step writes at lens[b]
+    page_tokens,
+    stride_kb, stride_kh, stride_vb, stride_vh,
+    stride_kp, stride_kph, stride_ks,
+    stride_vp, stride_vph, stride_vs,
+    stride_bt,
+    BLOCK_D: tl.constexpr,
+    D_K: tl.constexpr,
+    D_V: tl.constexpr,
+):
+    """A decode step's append with no host input: the position is read from the layer's DEVICE length and its page from
+    the block table, so the launch's arguments are the same at every step (a recorded replay replays it unchanged). One
+    program per (sequence, KV head); `paged_lens_advance_kernel`, launched after it on the same stream, moves the
+    length."""
+    pid_b = tl.program_id(0).to(tl.int64)
+    pid_h = tl.program_id(1).to(tl.int64)
+    pos = tl.load(lens_ptr + pid_b).to(tl.int64)
+    page = tl.load(block_table_ptr + pid_b * stride_bt + pos // page_tokens).to(tl.int64)
+    slot = pos % page_tokens
+    offs = tl.arange(0, BLOCK_D)
+    mk = offs < D_K
+    mv = offs < D_V
+    k = tl.load(k_new_ptr + pid_b * stride_kb + pid_h * stride_kh + offs, mask=mk, other=0.0)
+    v = tl.load(v_new_ptr + pid_b * stride_vb + pid_h * stride_vh + offs, mask=mv, other=0.0)
+    tl.store(k_pool_ptr + page * stride_kp + pid_h * stride_kph + slot * stride_ks + offs,
+             k.to(k_pool_ptr.dtype.element_ty), mask=mk)
+    tl.store(v_pool_ptr + page * stride_vp + pid_h * stride_vph + slot * stride_vs + offs,
+             v.to(v_pool_ptr.dtype.element_ty), mask=mv)
+
+
+@triton.jit
+def paged_lens_advance_kernel(lens_ptr, n):
+    """lens[b] += n for every sequence: one program per sequence, after the step's append (stream order)."""
+    b = tl.program_id(0).to(tl.int64)
+    tl.store(lens_ptr + b, tl.load(lens_ptr + b) + n)

nbxtensor/ops/pixel_shuffle.py

@@ -0,0 +1,144 @@
+"""Pixel shuffle (sub-pixel convolution) — pure @triton.jit kernel.
+
+Rearranges: [N, C*r*r, H, W] -> [N, C, H*r, W*r]
+
+For each output element (n, c, oh, ow):
+    ih = oh // r
+    iw = ow // r
+    ic = c * (r * r) + (oh % r) * r + (ow % r)
+    output[n, c, oh, ow] = input[n, ic, ih, iw]
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def pixel_shuffle_kernel(
+    input_ptr,
+    output_ptr,
+    n_elements,
+    # Input shape: [N, C_in, H, W] where C_in = C_out * r * r
+    C_out,
+    H,
+    W,
+    r,
+    # Output shape: [N, C_out, H*r, W*r]
+    OH,  # = H * r
+    OW,  # = W * r
+    # Input strides (contiguous NCHW)
+    stride_in_n,
+    stride_in_c,
+    stride_in_h,
+    stride_in_w,
+    # Output strides (contiguous NCHW)
+    stride_out_n,
+    stride_out_c,
+    stride_out_h,
+    stride_out_w,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Pixel shuffle: [N, C*r*r, H, W] -> [N, C, H*r, W*r].
+
+    Each thread handles one output element.
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    idx = 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 = idx < n_elements
+
+    # Decompose flat output index -> (n, c, oh, ow)
+    ow = idx % OW
+    tmp = idx // OW
+    oh = tmp % OH
+    tmp2 = tmp // OH
+    c = tmp2 % C_out
+    n = tmp2 // C_out
+
+    # Map output coords to input coords
+    ih = oh // r
+    iw = ow // r
+    # Sub-pixel position within the r x r block
+    sub_h = oh % r
+    sub_w = ow % r
+    ic = c * (r * r) + sub_h * r + sub_w
+
+    # Load from input, store to output
+    in_offset = n * stride_in_n + ic * stride_in_c + ih * stride_in_h + iw * stride_in_w
+    out_offset = n * stride_out_n + c * stride_out_c + oh * stride_out_h + ow * stride_out_w
+
+    val = tl.load(input_ptr + in_offset, mask=mask)
+    tl.store(output_ptr + out_offset, val, mask=mask)
+
+
+@triton.jit
+def pixel_shuffle_broadcast_aware_kernel(
+    bcast_view_ptr,  # 5D NBX expand-output view (N, C_pre, B, H, W) with stride_b == 0
+    output_ptr,      # 4D contiguous output (N, C_out, OH, OW)
+    n_elements,
+    C_out,           # output channels = (C_pre * B) // (r*r)
+    OH,              # = H * r
+    OW,              # = W * r
+    r,               # pixel_shuffle upscale factor
+    bcast,           # broadcast factor (= B, the size on the stride-0 dim)
+    # 5D bcast_view strides (NBX expand sets stride_b = 0).
+    stride_v_n,
+    stride_v_c,
+    stride_v_b,
+    stride_v_h,
+    stride_v_w,
+    # 4D output strides (contiguous NCHW).
+    stride_out_n,
+    stride_out_c,
+    stride_out_h,
+    stride_out_w,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Pixel shuffle that reads through an unmaterialized NBX broadcast view.
+
+    Input layout: (N, C_pre, B, H, W) where the original DAG pattern is
+    `unsqueeze(dim=2) -> expand(dim=2, B) -> clone -> view -> pixel_shuffle`.
+    NBX expand sets stride_b=0 (kernels/nbx_tensor.py:1581), so the 5D view
+    aliases the pre-expand tensor with no materialization. We read through
+    that view directly and skip the clone's 8 GiB copy.
+
+    The post-clone-post-view shape consumed by a normal pixel_shuffle is
+    (N, C_pre*B, H, W); here the kernel does the equivalent decomposition
+    inline:
+      ic_view = c*r*r + (oh%r)*r + (ow%r)         in [0, C_pre*B)
+      c_pre   = ic_view // bcast                   in [0, C_pre)
+      b       = ic_view %  bcast                   in [0, B)
+      ih      = oh // r
+      iw      = ow // r
+    Read offset uses all 5D strides; with stride_v_b == 0, the `b` index
+    contributes nothing and the load aliases naturally.
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    # Cast to int64 to support tensors with >= 2^31 elements
+    # (e.g. Sana 4Kpx VAE 1*128*4096*4096 = 2^31 — at INT32_MAX
+    # boundary; signed int32 multiplication of strides overflows
+    # silently and corrupts offsets, producing garbage output).
+    idx = 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 = idx < n_elements
+
+    ow = idx % OW
+    tmp = idx // OW
+    oh = tmp % OH
+    tmp2 = tmp // OH
+    c = tmp2 % C_out
+    n = tmp2 // C_out
+
+    sub_h = oh % r
+    sub_w = ow % r
+    ic_view = c * (r * r) + sub_h * r + sub_w
+    c_pre = ic_view // bcast
+    b = ic_view % bcast
+
+    ih = oh // r
+    iw = ow // r
+
+    in_offset = (n * stride_v_n + c_pre * stride_v_c + b * stride_v_b
+                 + ih * stride_v_h + iw * stride_v_w)
+    out_offset = (n * stride_out_n + c * stride_out_c
+                  + oh * stride_out_h + ow * stride_out_w)
+
+    val = tl.load(bcast_view_ptr + in_offset, mask=mask)
+    tl.store(output_ptr + out_offset, val, mask=mask)

nbxtensor/ops/pixel_unshuffle.py

@@ -0,0 +1,72 @@
+"""Pixel unshuffle (inverse of pixel shuffle) — pure @triton.jit kernel.
+
+Rearranges: [N, C, H*r, W*r] -> [N, C*r*r, H, W]
+
+For each output element (n, oc, oh, ow):
+    c = oc // (r * r)
+    sub_h = (oc % (r * r)) // r
+    sub_w = (oc % (r * r)) % r
+    ih = oh * r + sub_h
+    iw = ow * r + sub_w
+    output[n, oc, oh, ow] = input[n, c, ih, iw]
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def pixel_unshuffle_kernel(
+    input_ptr,
+    output_ptr,
+    n_elements,
+    # Input shape: [N, C, IH, IW] where IH = H*r, IW = W*r
+    C,
+    IH,
+    IW,
+    r,
+    # Output shape: [N, C*r*r, H, W]
+    C_out,  # = C * r * r
+    H,      # = IH // r
+    W,      # = IW // r
+    # Input strides (contiguous NCHW)
+    stride_in_n,
+    stride_in_c,
+    stride_in_h,
+    stride_in_w,
+    # Output strides (contiguous NCHW)
+    stride_out_n,
+    stride_out_c,
+    stride_out_h,
+    stride_out_w,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Pixel unshuffle: [N, C, H*r, W*r] -> [N, C*r*r, H, W].
+
+    Each thread handles one output element.
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    idx = 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 = idx < n_elements
+
+    # Decompose flat output index -> (n, oc, oh, ow)
+    ow = idx % W
+    tmp = idx // W
+    oh = tmp % H
+    tmp2 = tmp // H
+    oc = tmp2 % C_out
+    n = tmp2 // C_out
+
+    # Map output coords to input coords
+    c = oc // (r * r)
+    sub_idx = oc % (r * r)
+    sub_h = sub_idx // r
+    sub_w = sub_idx % r
+    ih = oh * r + sub_h
+    iw = ow * r + sub_w
+
+    # Load from input, store to output
+    in_offset = n * stride_in_n + c * stride_in_c + ih * stride_in_h + iw * stride_in_w
+    out_offset = n * stride_out_n + oc * stride_out_c + oh * stride_out_h + ow * stride_out_w
+
+    val = tl.load(input_ptr + in_offset, mask=mask)
+    tl.store(output_ptr + out_offset, val, mask=mask)

nbxtensor/ops/pow.py

@@ -0,0 +1,137 @@
+"""Power — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def pow_forward_kernel(
+    input_ptr, output_ptr,
+    n_elements,
+    exponent,
+    BLOCK_SIZE: tl.constexpr,
+    INT_EXP: tl.constexpr = 0,
+):
+    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)
+    x_fp32 = x.to(tl.float32)
+    # `tl.extra.CUDA.libdevice.pow` used to stand here. It is NVIDIA's device
+    # library, and asking any other backend to lower it fails at compile time —
+    # measured on triton-ext, Kokoro-82M dying at aten.pow::0 (1, 55, 2048) with
+    # "PassManager::run failed", and reproduced with a bare kernel containing no
+    # engine code (repro_pow_libdevice_cuda_namespace.py):
+    #     tl.extra.cuda.libdevice.pow   FAILED   PassManager::run failed
+    #     portable exp/log/sign         COMPILED max rel err 2.522e-07
+    # Triton 3.8.0 exposes no backend-neutral `pow`, so the semantics are built
+    # here rather than borrowed from one vendor.
+    #
+    # libm semantics, which is what the libdevice call was chosen for:
+    #   x > 0                     exp(e * ln x)
+    #   x < 0, e an integer       sign * exp(e * ln|x|), sign = -1 when e is odd
+    #   x < 0, e not an integer   NaN          (no real root)
+    #   x = +-0                   the THREE cases below, not one
+    #
+    # The zero lane read `0.0` for every exponent until 2026-09-19. That is
+    # wrong for two thirds of it, and the Dell measured it on CUDA against an
+    # fp64 oracle: `0.0 ** 0.0` gave 0.0 where the oracle gives 1.0, and
+    # `0.0 ** -1.0` gave 0.0 where it gives inf. Mine, and the shape of the
+    # mistake is worth naming: replacing a vendor intrinsic means reimplementing
+    # its whole DOMAIN, and I reimplemented the interesting part and not the
+    # edges. C99 / IEEE-754 pow at zero:
+    #
+    #   pow(+-0, 0)                1              (even 0**0, and even nan**0)
+    #   pow(+-0, e) e > 0          +-0, negative only for -0 with e an odd int
+    #   pow(+-0, e) e < 0          +-inf, same sign rule
+    #
+    # The sign of the zero has to be read from the SIGN BIT, because `x < 0` is
+    # false for -0.0 and `x == 0` cannot tell the two apart. A bitcast to int32
+    # is negative exactly when the sign bit is set, which covers -0.0.
+    _ax = tl.abs(x_fp32)
+    _mag = tl.exp(exponent * tl.log(_ax))
+    _is_int = tl.floor(exponent) == exponent
+    _is_odd = _is_int & ((tl.floor(exponent * 0.5) * 2.0) != exponent)
+    _sign = tl.where(_is_odd, -1.0, 1.0)
+    _neg = tl.where(_is_int, _sign * _mag, float("nan"))
+
+    _signbit = x_fp32.to(tl.int32, bitcast=True) < 0      # true for -0.0 too
+    _zsign = tl.where(_signbit & _is_odd, -1.0, 1.0)
+    _zero = tl.where(exponent == 0.0, 1.0,
+                     tl.where(exponent > 0.0,
+                              _zsign * 0.0,               # +-0
+                              _zsign * float("inf")))     # +-inf
+    out = tl.where(x_fp32 > 0, _mag, tl.where(x_fp32 < 0, _neg, _zero))
+    # THE EXACT ROUTE FOR A SMALL INTEGER EXPONENT — the branch the comment
+    # below records as "available rather than taken". Taken on 2026-09-20,
+    # because the Dell measured the general route on CUDA against the fp64
+    # oracle and it is not a rounding footnote: x**2 lands up to 15 ulps from
+    # the oracle where main's libdevice call sat within 2 (38 147 lanes, |x| a
+    # standard normal with zero lanes), x**0.5 at 5 vs 2, x**-1 at 7 vs 2. A
+    # square computed as exp(2 log|x|) is not a square. Repeated multiplication
+    # is correctly rounded at e=2 and within two roundings at e=3, needs no
+    # vendor library, and carries IEEE's own sign rules for free: (-0)**2 is
+    # +0, (-0)**3 is -0, 1/(+-0) is +-inf, so none of the select logic above is
+    # needed on this route. The wrapper sets INT_EXP when the exponent is an
+    # integer-valued scalar with 1 <= |e| <= 8; INT_EXP == 0 keeps the general
+    # route, which still owns e == 0, non-integers and large integers.
+    if INT_EXP != 0:
+        _n: tl.constexpr = INT_EXP if INT_EXP > 0 else -INT_EXP
+        _p = x_fp32
+        if _n >= 2:
+            _p = _p * x_fp32
+        if _n >= 3:
+            _p = _p * x_fp32
+        if _n >= 4:
+            _p = _p * x_fp32
+        if _n >= 5:
+            _p = _p * x_fp32
+        if _n >= 6:
+            _p = _p * x_fp32
+        if _n >= 7:
+            _p = _p * x_fp32
+        if _n >= 8:
+            _p = _p * x_fp32
+        if INT_EXP < 0:
+            _p = 1.0 / _p
+        out = _p
+    # ---------------------------------------------------------------------
+    # THE ACCURACY COST OF THIS ROUTE, AND WHAT IT IS AND IS NOT INHERENT TO.
+    #
+    # `exp(e * log|x|)` loses precision as |x| moves away from 1, because the
+    # error in `log|x|` is multiplied by `e` before `exp` amplifies it. The Dell
+    # measured pow-to-the-third at 19x worse relative error than main's
+    # libdevice call on CUDA. Measured here against fp64, on the ranges a
+    # vocoder actually sees (probe_pow_error.py):
+    #
+    #   range                    e    this kernel   x*x(*x)     ratio
+    #   |x| in [0.5, 2]          2      2.56e-07   5.94e-08      4.3x
+    #   |x| in [0.5, 2]          3      3.95e-07   1.14e-07      3.4x
+    #   |x| in [1e-3, 1e3]       3      3.58e-06   1.15e-07       31x
+    #   |x| in [1e-6, 1e-3]      3      7.24e-06   1.13e-07       64x
+    #
+    # So the 19x is real and it is range-dependent: 3.4x on typical activations,
+    # 64x on small ones.
+    #
+    # WHAT IT COSTS ON THE MODELS THAT USE IT. `aten::pow` appears 194 times
+    # across the 16 cached containers, and only in the two vocoder TTS models:
+    #   exponent 2    182x   Kokoro-82M 49, chatterbox 133
+    #   exponent 3.0   12x   Kokoro-82M only
+    # Nothing else in the catalogue calls pow at all, and no site uses a
+    # non-integer exponent. Both models are judged by third-party ASR and both
+    # transcribe correctly; the error is ~1e-7 relative on their activation
+    # range, which is below fp32 round-off for the accumulations downstream of
+    # it. So the cost TODAY is not audible and not measurable at the artefact.
+    #
+    # BUT IT IS AVOIDABLE, and that is the part worth recording rather than
+    # accepting: the cost is inherent to `exp(e * log|x|)`, NOT to `pow`. Every
+    # exponent the catalogue uses is a small integer, and for those, repeated
+    # multiplication is 3-64x MORE accurate and no more expensive — one
+    # multiply for e=2 against a log, a multiply and an exp. A `POW_SMALL_INT`
+    # constexpr branch would take the exact route at compile time and leave this
+    # path for the general case.
+    #
+    # It is NOT done here, deliberately: it changes the numerics of two SHIPPED
+    # models, so it owes a fresh set of judged cells and a fresh CUDA column,
+    # and 0.5.4 has just been cut. Recorded as available rather than taken.
+    tl.store(output_ptr + offset, out, mask=mask)

nbxtensor/ops/prod.py

@@ -0,0 +1,81 @@
+"""Product reduction — pure @triton.jit kernels.
+
+Global reduction via two-pass (mid → result) and dim-specific reduction.
+Extracted from FlagGems (Apache-2.0).
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def reduce_mul(a, b):
+    return a * b
+
+
+# ---------------------------------------------------------------------------
+# Global reduction: two-pass partial products
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def prod_kernel_mid(
+    inp,
+    mid,
+    M,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Pass 1: each program computes a partial product over its BLOCK_SIZE chunk."""
+    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)
+    inp_ptrs = inp + offset
+    mask = offset < M
+    inp_val = tl.load(inp_ptrs, mask=mask, other=1.0).to(tl.float32)
+    mid_value = tl.reduce(inp_val, axis=0, combine_fn=reduce_mul)
+    mid_ptr = mid + pid
+    tl.store(mid_ptr, mid_value.to(inp_val.dtype))
+
+
+@triton.jit
+def prod_kernel_result(mid, out, mid_size, BLOCK_MID: tl.constexpr):
+    """Pass 2: reduce all partial products into a single scalar."""
+    offset = tl.arange(0, BLOCK_MID)
+    mid_ptrs = mid + offset
+    mask = offset < mid_size
+    mid_val = tl.load(mid_ptrs, mask=mask, other=1.0).to(tl.float32)
+    prod_val = tl.reduce(mid_val, axis=0, combine_fn=reduce_mul)
+    tl.store(out, prod_val)
+
+
+# ---------------------------------------------------------------------------
+# Dim-specific reduction: product along one axis
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def prod_kernel(
+    inp,
+    out,
+    M,
+    N,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """Reduce product along the compressed inner dimension N.
+
+    inp: [M, N] (dim-compressed) → out: [M]
+    """
+    pid_m = tl.program_id(0).to(tl.int64)
+    m_offset = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+
+    acc = tl.full((BLOCK_M, BLOCK_N), value=1.0, dtype=tl.float32)
+    for start_n in range(0, N, BLOCK_N):
+        n_offset = start_n + tl.arange(0, BLOCK_N)
+        offset = m_offset[:, None] * N + n_offset[None, :]
+        mask = (m_offset[:, None] < M) & (n_offset[None, :] < N)
+        inp_ptrs = inp + offset
+        inp_vals = tl.load(inp_ptrs, mask=mask, other=1.0).to(tl.float32)
+        acc *= inp_vals
+    result = tl.reduce(acc, axis=1, combine_fn=reduce_mul)
+
+    out_ptrs = out + m_offset
+    mask1 = m_offset < M
+    tl.store(out_ptrs, result, mask=mask1)