nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/decode_attn_vec.py

@@ -0,0 +1,162 @@
+"""Vector (SIMT) decode attention — Tq == 1, pure @triton.jit, no tl.dot.
+
+Second attempt at the decode-attention bar (~0.5 ms/token at 4k = the
+KV bytes at HBM speed). The first attempt (ops/flash_decode.py) kept
+the flash-decoding two-pass SPLIT structure but computed QK/PV with
+`tl.dot` on a 16-row Q tile carrying 4 real GQA rows — and lost on the
+row three times. The 2026-08-23 web research (R16) showed the split
+was never the problem and the dot was: every production Tq=1 kernel on
+pre-Ampere silicon is a SIMT kernel — FasterTransformer's
+masked_multihead_attention (per-thread fp16 FMA + shuffle reduce, no
+MMA anywhere), vLLM's paged_attention v1/v2 (16-byte vectorized K
+loads, fixed-order v2 partition reduce), llama.cpp's fattn-vec/tile
+(picked over MMA on Volta at exactly our effective batch), lightllm's
+token_attention (`tl.sum(q[None, :] * k, 1)`, BLOCK_N=32, no dot).
+Triton lowers fp16 `tl.dot` to FFMA on sm_70 (no HMMA lowering — the
+Phase 1.5 measurement), so the MMA tile staging is pure cost there.
+
+Structure here (the lightllm/vLLM shape):
+  - grid = (B * H_q, split): ONE QUERY ROW per program — no G tile at
+    all. The GQA K re-read (4 Q heads share a KV head) is absorbed by
+    L2 temporal locality, exactly as in vLLM/FT which organise the
+    same way.
+  - QK: `tl.sum(q[None, :] * k_tile, 1)` — an FMA reduction over D,
+    product tile (BLOCK_N, BLOCK_D) only.
+  - online softmax with SCALAR m/l state, fp32 accumulator [BLOCK_D].
+  - split-KV partials (out, m, l) merged by the EXISTING fixed-order
+    `flash_decode_reduce_kernel` (ops/flash_decode.py) viewed at
+    GROUPS=1 — deterministic by construction, no atomics, proven.
+  - K/V are read by their OWN strides: the bucketed KV cache's
+    prefix-slice views flow in directly (the strided-KV lot's
+    contract), no materialisation.
+  - num_stages=1 (sm_70 has no async-copy engine to pipeline with).
+
+Born with its float64 oracle at D = 127, 128 AND 129
+(test_decode_vec_oracle.py) and a route ACTIVATION proof — the
+2026-08-23 lesson: a route-dependent claim needs the route pinned.
+"""
+
+import triton
+import triton.language as tl
+
+from nbxtensor.ops._configs import nbx_autotune
+
+
+@triton.jit
+def _round_q(q, Q_TO: tl.constexpr, Q_SATURATE: tl.constexpr):
+    """Q as the caller asked it read: rounded to Q_TO (the KV cache's
+    dtype) with the stored cast's protected fp16 conversion when
+    Q_SATURATE, then widened to the fp32 the dot runs in — the bytes of a
+    cast copy of Q, without the copy."""
+    if Q_TO is not None:
+        f = q.to(tl.float32)
+        if Q_SATURATE:
+            is_finite = (f == f) & (f != float("inf")) & (f != float("-inf"))
+            clamped = tl.minimum(tl.maximum(f, -65504.0), 65504.0)
+            f = tl.where(is_finite, clamped, f)
+        q = f.to(Q_TO)
+    return q.to(tl.float32)
+
+
+#: the split: one program per DECODE_VEC_SEG keys (the judged split, locked pinned protocol, 2026-08-23: 256 won the
+#: row at all three contexts), at most DECODE_VEC_MAX_SPLIT — it defines the partials, so it is the algorithm's, not a
+#: tile
+DECODE_VEC_SEG = 256
+DECODE_VEC_MAX_SPLIT = 64
+
+
+def decode_vec_splits(seqlen_k: int) -> int:
+    """The splits the vector decode cuts `seqlen_k` keys into — its grid, and the shared reduce's autotune key (one
+    rule for the wrapper, the census and the certifier)."""
+    return max(1, min(-(-int(seqlen_k) // DECODE_VEC_SEG), DECODE_VEC_MAX_SPLIT))
+
+
+# The key block and warps of the vector decode, candidates the certification measures per key (the judged row,
+# 2026-08-23: 32 keys a step, 4 warps, among them); one stage — the loop streams K and V once.
+_DECODE_VEC_CONFIGS = [triton.Config({'BLOCK_N': bn}, num_warps=w, num_stages=1)
+                       for bn in (16, 32, 64, 128) for w in (1, 2, 4, 8)]
+
+
+@nbx_autotune(configs=_DECODE_VEC_CONFIGS, key=['T_BUCKET', 'D', 'D_V', 'GQA_GROUPS', 'HAS_BIAS'], cache_results=True)
+@triton.jit
+def decode_attn_vec_split_kernel(
+    q_ptr,          # [B*H_q, D]        one query row per program (fp16)
+    k_ptr,          # [B*H_kv, T_k, D]  native cache layout, own strides
+    v_ptr,          # [B*H_kv, T_k, D_v]
+    bias_ptr,       # [T_k] additive fp32/fp16 (unused when HAS_BIAS=0)
+    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
+    T_k, seg_len,
+    T_BUCKET,       # T_k's bucket: the autotune key's length class (unused in the body)
+    sm_scale,
+    stride_qh,
+    stride_kh, stride_kn,
+    stride_vh, stride_vn,
+    stride_oh, stride_os,
+    stride_mh,
+    GQA_GROUPS: tl.constexpr,   # H_q // H_kv — maps Q head -> KV head
+    BLOCK_D: tl.constexpr,
+    D: tl.constexpr,
+    D_V: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    Q_TO: tl.constexpr = None,        # a dtype: Q is rounded to it on load (the cache's)
+    Q_SATURATE: tl.constexpr = False, # the protected fp16 rounding of that cast
+    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
+    pid_kv = pid_q // GQA_GROUPS    # owning KV head row
+
+    offs_d = tl.arange(0, BLOCK_D)
+    mask_d = offs_d < D
+    mask_dv = offs_d < D_V
+
+    q = tl.load(q_ptr + pid_q * stride_qh + offs_d,
+                mask=mask_d, other=0.0)
+    q = _round_q(q, Q_TO, Q_SATURATE)
+
+    seg_start = pid_s * seg_len
+    seg_start = seg_start.to(tl.int32)   # loop bounds are 32-bit counts: the Metal lowering refuses a 64-bit scf.for bound (addressing stays 64-bit through the program ids)
+    seg_end = tl.minimum(seg_start + seg_len, T_k).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):
+        offs_n = n0 + tl.arange(0, BLOCK_N)
+        mask_n = offs_n < seg_end
+
+        k_tile = tl.load(k_ptr + pid_kv * stride_kh
+                         + offs_n[:, None] * stride_kn + offs_d[None, :],
+                         mask=mask_n[:, None] & mask_d[None, :], other=0.0)
+        # scores [BLOCK_N] — FMA reduction over D, no tl.dot
+        s = tl.sum(q[None, :] * k_tile.to(tl.float32), 1) * sm_scale
+        if HAS_BIAS:
+            b = tl.load(bias_ptr + offs_n, mask=mask_n, other=0.0)
+            s = s + b.to(tl.float32)
+        s = tl.where(mask_n, s, float("-inf"))
+
+        m_new = tl.maximum(m_i, tl.max(s, 0))
+        # empty tail keeps m_new=-inf; pin the exp arguments (the
+        # flash_decode kernel's own guard, scalar form)
+        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_ptr + pid_kv * stride_vh
+                         + offs_n[:, None] * stride_vn + 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)
+
+

nbxtensor/ops/depthwise_conv2d.py

@@ -0,0 +1,132 @@
+"""Depthwise conv2d — direct stencil @triton.jit kernel.
+
+Routed by `kernels/wrappers.py:conv2d_wrapper` when the convolution
+matches the depthwise signature (`groups == in_c == out_c`, weight
+shape `(C, 1, kh, kw)`). The generic `conv2d_forward_kernel` (im2col +
+tl.dot) is structurally inefficient for this case: each group has K =
+kh*kw (often 9) so the inner reduction loop runs once with most of the
+BLOCK_INF lane idle, yet the kernel still pays full launch + im2col
+cost for every group. On Sana 4Kpx VAE (the DC-AE depthwise blocks
+groups=in=out=11200), the generic kernel measures ~4.8 s per call
+versus ~2.6 ms for cuDNN's dedicated depthwise path — a ~1800× gap
+that closes under the dedicated stencil.
+
+Pattern reference (P-SANA-4KPX-RUNTIME Step 3 study):
+- MultiPath/DepthwiseConv2d (CUTLASS sm_70 NCHW depthwise iterator,
+  threadblock 64x128x32, warp 32x32x32, instruction 8x8x4 HMMA)
+- PyTorch PR #22302 — confirmed cuDNN has a separate depthwise code
+  path on Volta/Turing fp16 (cudnn >= 7600)
+
+Our Triton-pure design: each program computes a (BC channels x TH rows x TW columns) output tile of one batch
+sample, width the fastest axis (NCHW's contiguous one), accumulating the kh*kw stencil contributions in fp32: the
+kh*kw shifted input windows of a compact 2-D tile overlap and are served from the L1. No im2col buffer, no
+cross-channel reduction, no HMMA pseudo-GEMM packing; the stencil is unrolled at compile time (kh, kw constexpr).
+The tile was (BLOCK_HW flattened positions x BLOCK_C channels) with the channels fastest, H*W elements apart in
+memory: no load coalesced — SANA-Video's 11200-channel 3x3 fp32 [2, 11200, 24, 40] ran 2.63 ms where this one runs
+0.33 ms (529 GB/s; torch eager 0.38 ms), [22, 13440, 21, 42] 18.9 -> 4.2 ms (V100-SXM2, 2026-11-03).
+"""
+
+import triton
+
+from nbxtensor.ops._configs import nbx_autotune
+import triton.language as tl
+
+from nbxtensor.ops._autotune_policy import maybe_pin_single, is_depthwise_conv2d_pinned
+
+
+_DEPTHWISE_CONV2D_CONFIGS = [
+    triton.Config({'BC': 1, 'TH': 8, 'TW': 64}, num_stages=2, num_warps=4),
+    triton.Config({'BC': 1, 'TH': 16, 'TW': 64}, num_stages=2, num_warps=4),
+    triton.Config({'BC': 1, 'TH': 32, 'TW': 64}, num_stages=2, num_warps=8),
+    triton.Config({'BC': 2, 'TH': 8, 'TW': 64}, num_stages=2, num_warps=4),
+    triton.Config({'BC': 4, 'TH': 8, 'TW': 32}, num_stages=2, num_warps=4),
+    triton.Config({'BC': 4, 'TH': 8, 'TW': 64}, num_stages=2, num_warps=8),
+    triton.Config({'BC': 8, 'TH': 4, 'TW': 64}, num_stages=2, num_warps=8),
+    triton.Config({'BC': 16, 'TH': 4, 'TW': 32}, num_stages=2, num_warps=8),
+    triton.Config({'BC': 1, 'TH': 4, 'TW': 128}, num_stages=2, num_warps=4),
+    # one-row tiles: the census's depthwise rows are mostly conv1d run at H = 1 (1 x 7, 64-1024 channels), where a
+    # taller tile is all masked rows — measured best there (1024 x 2048 fp16: 0.027 ms, the flattened tile 0.062)
+    triton.Config({'BC': 4, 'TH': 1, 'TW': 128}, num_stages=2, num_warps=4),
+    triton.Config({'BC': 8, 'TH': 1, 'TW': 64}, num_stages=2, num_warps=4),
+    triton.Config({'BC': 16, 'TH': 1, 'TW': 32}, num_stages=2, num_warps=4),
+    triton.Config({'BC': 16, 'TH': 1, 'TW': 64}, num_stages=2, num_warps=8),
+    triton.Config({'BC': 32, 'TH': 1, 'TW': 32}, num_stages=2, num_warps=8),
+]
+
+
+@nbx_autotune(
+    configs=maybe_pin_single(_DEPTHWISE_CONV2D_CONFIGS, is_depthwise_conv2d_pinned),
+    key=['C', 'H_in', 'W_in_key', 'H_out', 'W_out_key', 'kh', 'kw',
+         'stride_h', 'stride_w', 'pad_h', 'pad_w', 'fp16'],
+    cache_results=True,
+)
+@triton.jit
+def depthwise_conv2d_kernel(
+    x_ptr, w_ptr, out_ptr,
+    N, C,
+    H_in, W_in, H_out, W_out,
+    W_in_key, W_out_key,
+    x_n_stride, x_c_stride, x_h_stride, x_w_stride,
+    w_c_stride, w_kh_stride, w_kw_stride,
+    out_n_stride, out_c_stride, out_h_stride, out_w_stride,
+    kh: tl.constexpr, kw: tl.constexpr,
+    stride_h: tl.constexpr, stride_w: tl.constexpr,
+    pad_h: tl.constexpr, pad_w: tl.constexpr,
+    fp16: tl.constexpr,
+    BC: tl.constexpr = 2,
+    TH: tl.constexpr = 8,
+    TW: tl.constexpr = 64,
+    WIDE_OFFSETS: tl.constexpr = True,
+):
+    # 64-bit offsets only where an operand's extent needs them (`WIDE_OFFSETS`, the wrapper's `_conv_offsets_wide`, as
+    # conv2d): 32-bit offsets are exact below 2^31, and the 64-bit ones are emulated on Apple's ALUs — under the i64 fix of
+    # the Metal backend (b06cbb0d) the fp32 3x3 depthwise (13440 ch, 42x21) ran 8.4 -> 12.5 ms (Mac, 2026-10-21)
+    if WIDE_OFFSETS:
+        pid_nc = tl.program_id(0).to(tl.int64)
+        pid_h = tl.program_id(1).to(tl.int64)
+        pid_w = tl.program_id(2).to(tl.int64)
+    else:
+        pid_nc = tl.program_id(0)
+        pid_h = tl.program_id(1)
+        pid_w = tl.program_id(2)
+    n_cb = tl.cdiv(C, BC)
+    pid_n = pid_nc // n_cb
+    offs_c = (pid_nc % n_cb) * BC + tl.arange(0, BC)
+    oh = pid_h * TH + tl.arange(0, TH)
+    ow = pid_w * TW + tl.arange(0, TW)
+    c_mask = offs_c < C
+
+    accum = tl.zeros((BC, TH, TW), dtype=tl.float32)
+    x_base = x_ptr + pid_n * x_n_stride + offs_c[:, None, None] * x_c_stride
+
+    # Unrolled stencil over kh x kw (constexpr): each tap reads the input window the output tile shifted by it.
+    for kh_i in tl.static_range(0, kh):
+        ih = oh * stride_h + kh_i - pad_h
+        h_ok = (ih >= 0) & (ih < H_in)
+        for kw_i in tl.static_range(0, kw):
+            iw = ow * stride_w + kw_i - pad_w
+            w_ok = (iw >= 0) & (iw < W_in)
+            full_mask = c_mask[:, None, None] & h_ok[None, :, None] & w_ok[None, None, :]
+            x_block = tl.load(x_base + ih[None, :, None] * x_h_stride + iw[None, None, :] * x_w_stride,
+                              mask=full_mask, other=0.0)
+            w_block = tl.load(w_ptr + offs_c * w_c_stride + kh_i * w_kh_stride + kw_i * w_kw_stride,
+                              mask=c_mask, other=0.0)
+            # Upcast for EVERY dtype, not just fp16. The accumulator is fp32 either way, so
+            # this costs nothing — and multiplying in the operands' own dtype made bf16 with
+            # padding wrong on Metal: deviation 0.754 against the fp64 oracle where unpadded
+            # bf16 sat at its 3.2e-03 mantissa floor, with the error in the ALL-IN-BOUNDS
+            # interior while the genuinely masked border row was clean. fp32 and fp16 were
+            # unaffected at both paddings, and the same source is exact to the mantissa on
+            # CUDA (Triton 3.8.0, V100 sm_70: bf16 pad0 == pad1), so the bad lowering is the
+            # Metal backend's for a native bf16 product of a masked-loaded operand. Upcasting
+            # removes the dependence on it rather than relying on it being fixed below us.
+            # Gate: tests/unit/kernels/test_depthwise_bf16_padding.py (2026-09-22).
+            accum += x_block.to(tl.float32) * w_block.to(tl.float32)[:, None, None]
+
+    out_addr = (out_ptr + pid_n * out_n_stride + offs_c[:, None, None] * out_c_stride
+                + oh[None, :, None] * out_h_stride + ow[None, None, :] * out_w_stride)
+    out_mask = c_mask[:, None, None] & (oh < H_out)[None, :, None] & (ow < W_out)[None, None, :]
+    if fp16:
+        tl.store(out_addr, accum.to(tl.float16), mask=out_mask)
+    else:
+        tl.store(out_addr, accum, mask=out_mask)

nbxtensor/ops/dequant_gemv.py

@@ -0,0 +1,190 @@
+"""E7 — fused int4 dequant-GEMV family (quantized tier, pure Triton).
+
+Weight-only int4, groupwise g=128, ASYMMETRIC: per group along K an
+fp16 scale and an fp16 min; dequant is `q * scale + min` with q an
+UNSIGNED nibble (llama.cpp q4_1's exact form — no -8 offset, no
+AutoGPTQ +1 storage quirk). Packing: 8 nibbles per int32, flat
+low-nibble-first (nibble i of word r encodes element k = 8*r + i),
+W_q laid out (K//8, N) K-major so loads coalesce along N — the
+GemLite convention.
+
+THE TIER'S BYTE ORACLE (scoping "quantized tier contract", clause 1):
+the fused kernel must be byte-identical to dequantize-then-GEMV. A
+cuBLAS/torch reference can never be byte-matched (its reduction order
+is internal); the oracle is therefore built from THIS file's own
+kernels: `dequant_int4_kernel` (standalone dequant whose dequant
+expression is textually identical to the fused kernel's) feeding
+`gemv_ref_kernel` (same BLOCK_N/BLOCK_K, same ascending-K chunk
+loop, same tl.sum tree, same fp32 accumulator). Fused == oracle by
+construction; the microtest proves it byte-for-byte with ONE pinned
+config for both kernels (autotune would change the tl.sum tree — E7
+GEMV sits OUTSIDE the mm/bmm/addmm/conv2d autotune sanction).
+
+sm_70 discipline (sourced, R38 study + 2026-08-16 research): M=1
+decode GEMV, no tl.dot (Volta fp16 dot lowers to FMA anyway), SPLIT_K
+= 1 — a single program per N-tile walks all of K, NO atomics
+(atomic-splitK is arrival-order nondeterministic and fp16 atomics are
+version-fragile on our pin; both break the byte contract). The win is
+weight bandwidth: 4x fewer weight bytes than fp16. num_stages <= 2
+(cp.async pipelining is sm_80+; more stages only add register
+pressure to the unpack-heavy loop).
+
+Constraints enforced by the wrappers: K % BLOCK_K == 0 handled by
+masking; BLOCK_K in {32, 64, 128} divides group_size=128 so a K-chunk
+never crosses a group boundary (one scale/min scalar per (chunk, n) —
+the GemLite pruner rule).
+"""
+
+import triton
+import triton.language as tl
+
+GROUP_SIZE: int = 128
+PACK: int = 8  # nibbles per int32
+
+# Pinned parity config: oracle AND fused always run the SAME values —
+# the tl.sum reduction tree must be identical between the two, which
+# is the tier contract (fused == dequant->ref byte-for-byte).
+# BLOCK_K and NUM_WARPS stay pinned. BLOCK_N adapts at runtime for
+# small-N grids (wrappers._dequant_gemv_block_n, 2026-08-25: halved
+# pow2 until the grid reaches the vendor-yml occupancy target,
+# NBX_DGEMV_ADAPT=0 kill) — legitimate because BLOCK_N only
+# partitions OUTPUT columns while each column's K-reduction tree is
+# BLOCK_K-fixed; the parity microtest proves fused == oracle
+# byte-for-byte at the adapted widths (1/8) as well as the pinned 64.
+# Cross-WIDTH outputs differ in low bits (the per-column lowering
+# changes with BLOCK_N) — width changes are drift-class and go
+# through the model byte gate, never silently.
+BLOCK_N: int = 64
+BLOCK_K: int = 128
+NUM_WARPS: int = 4
+
+
+@triton.jit
+def dequant_int4_kernel(
+    wq_ptr, scales_ptr, mins_ptr, out_ptr,
+    K, N,
+    stride_wk, stride_wn,
+    stride_sg, stride_sn,
+    stride_ok, stride_on,
+    BLOCK_K_C: tl.constexpr, BLOCK_N_C: tl.constexpr,
+    GROUP_C: tl.constexpr, PACK_C: tl.constexpr,
+):
+    """Standalone dequant: packed int4 -> dense fp32 [K, N].
+
+    Oracle part 1. The dequant expression below is the TEXTUAL
+    contract shared with the fused kernel: same loads, same shifts,
+    same pure-fp32 dtype path.
+    """
+    pid_k = tl.program_id(0).to(tl.int64)
+    pid_n = tl.program_id(1).to(tl.int64)
+    offs_k = pid_k * BLOCK_K_C + tl.arange(0, BLOCK_K_C)
+    offs_n = pid_n * BLOCK_N_C + tl.arange(0, BLOCK_N_C)
+    mask_k = offs_k < K
+    mask_n = offs_n < N
+    mask = mask_k[:, None] & mask_n[None, :]
+
+    packed = tl.load(
+        wq_ptr + (offs_k[:, None] // PACK_C) * stride_wk
+        + offs_n[None, :] * stride_wn, mask=mask, other=0)
+    q = (packed >> ((offs_k[:, None] % PACK_C) * 4)) & 0xF
+
+    g = pid_k * BLOCK_K_C // GROUP_C  # BLOCK_K_C divides GROUP_C
+    scale = tl.load(scales_ptr + g * stride_sg + offs_n * stride_sn,
+                    mask=mask_n, other=0.0)
+    mn = tl.load(mins_ptr + g * stride_sg + offs_n * stride_sn,
+                 mask=mask_n, other=0.0)
+    # Canonical dequant: PURE fp32 (llama.cpp dmmv default,
+    # dfloat=float — the fp16 weight never exists at runtime).
+    # Contraction-immune: q*scale is EXACT in fp32 (4-bit int x fp16
+    # mantissa), so fma(q,s,m) == q*s+m bit-for-bit whatever the
+    # compiler emits. The dense oracle tensor is fp32.
+    w = (q.to(tl.float32) * scale[None, :].to(tl.float32)
+         + mn[None, :].to(tl.float32))
+
+    tl.store(out_ptr + offs_k[:, None] * stride_ok
+             + offs_n[None, :] * stride_on, w, mask=mask)
+
+
+@triton.jit
+def gemv_ref_kernel(
+    x_ptr, w_ptr, out_ptr,
+    K, N,
+    stride_wk, stride_wn,
+    BLOCK_K_C: tl.constexpr, BLOCK_N_C: tl.constexpr,
+):
+    """Reference dense-weight GEMV: out[n] = sum_k x[k] * w[k, n].
+
+    Oracle part 2. Fixed reduction structure — ascending-K chunk loop,
+    per-chunk tl.sum over axis 0, fp32 accumulator, one store — the
+    EXACT structure of the fused kernel below. (A nibble-major
+    restructure with per-nibble partial sums was MEASURED and
+    reverted: it slowed the fused kernel 0.090 -> 0.224 ms at
+    4096x11008 — the lane-selection sums cost more than the redundant
+    packed loads they save; the fused kernel is ALU-issue-bound, not
+    load-bound.)
+    """
+    pid_n = tl.program_id(0).to(tl.int64)
+    offs_n = pid_n * BLOCK_N_C + tl.arange(0, BLOCK_N_C)
+    mask_n = offs_n < N
+    acc = tl.zeros((BLOCK_N_C,), dtype=tl.float32)
+    for kc in range(0, tl.cdiv(K, BLOCK_K_C)):
+        offs_k = kc * BLOCK_K_C + tl.arange(0, BLOCK_K_C)
+        mask_k = offs_k < K
+        a = tl.load(x_ptr + offs_k, mask=mask_k, other=0.0)
+        b = tl.load(w_ptr + offs_k[:, None] * stride_wk
+                    + offs_n[None, :] * stride_wn,
+                    mask=mask_k[:, None] & mask_n[None, :], other=0.0)
+        acc += tl.sum(a.to(tl.float32)[:, None] * b.to(tl.float32),
+                      axis=0)
+    tl.store(out_ptr + offs_n, acc, mask=mask_n)
+
+
+@triton.jit
+def dequant_gemv_int4_kernel(
+    x_ptr, wq_ptr, scales_ptr, mins_ptr, out_ptr,
+    K, N,
+    stride_wk, stride_wn,
+    stride_sg, stride_sn,
+    BLOCK_K_C: tl.constexpr, BLOCK_N_C: tl.constexpr,
+    GROUP_C: tl.constexpr, PACK_C: tl.constexpr,
+):
+    """Fused: out[n] = sum_k x[k] * dequant(wq)[k, n], fp32 accum.
+
+    Dequant expression textually identical to dequant_int4_kernel;
+    reduction structure textually identical to gemv_ref_kernel —
+    byte-equal to the two-step oracle by construction (proven by the
+    parity microtest with the pinned config). The packed load gathers
+    `offs_k // 8` (each int32 requested 8x, absorbed by L1) — the
+    nibble-major single-load variant was MEASURED SLOWER (see the
+    oracle kernel's docstring); the kernel is ALU-issue-bound.
+    """
+    pid_n = tl.program_id(0).to(tl.int64)
+    offs_n = pid_n * BLOCK_N_C + tl.arange(0, BLOCK_N_C)
+    mask_n = offs_n < N
+    acc = tl.zeros((BLOCK_N_C,), dtype=tl.float32)
+    for kc in range(0, tl.cdiv(K, BLOCK_K_C)):
+        offs_k = kc * BLOCK_K_C + tl.arange(0, BLOCK_K_C)
+        mask_k = offs_k < K
+        mask = mask_k[:, None] & mask_n[None, :]
+        a = tl.load(x_ptr + offs_k, mask=mask_k, other=0.0)
+
+        packed = tl.load(
+            wq_ptr + (offs_k[:, None] // PACK_C) * stride_wk
+            + offs_n[None, :] * stride_wn, mask=mask, other=0)
+        q = (packed >> ((offs_k[:, None] % PACK_C) * 4)) & 0xF
+        g = kc * BLOCK_K_C // GROUP_C  # BLOCK_K_C divides GROUP_C
+        scale = tl.load(scales_ptr + g * stride_sg + offs_n * stride_sn,
+                        mask=mask_n, other=0.0)
+        mn = tl.load(mins_ptr + g * stride_sg + offs_n * stride_sn,
+                     mask=mask_n, other=0.0)
+        # Canonical dequant — textually identical to
+        # dequant_int4_kernel: PURE fp32, no fp16 materialization
+        # (contraction-immune: q*scale exact in fp32).
+        w = (q.to(tl.float32) * scale[None, :].to(tl.float32)
+             + mn[None, :].to(tl.float32))
+        # Masked K-tail rows dequant to `mn` (q=0), not 0 — zero them
+        # through the mask so the tail contributes exact +0.0, same as
+        # the oracle GEMV's masked b-load.
+        b = tl.where(mask, w, 0.0)
+        acc += tl.sum(a.to(tl.float32)[:, None] * b, axis=0)
+    tl.store(out_ptr + offs_n, acc, mask=mask_n)

nbxtensor/ops/dft.py

@@ -0,0 +1,67 @@
+"""DFT / iDFT basis matrices — pure @triton.jit kernels.
+
+Builds the real cos/sin basis matrices for a Discrete Fourier Transform computed
+by matrix multiplication. This is the standard small-N / non-power-of-2 path
+(the DFT is an N×N matrix of complex exponentials applied to the signal) — used
+when the radix-2 butterfly cannot run (n_fft not a power of 2, e.g. an iSTFT head
+with n_fft=20). Model-agnostic: any non-pow2 FFT length routes here.
+
+Layout (matches the rfft convention): frequency bins k in [0, N//2+1), samples
+n in [0, N). The matrices are [N_bins, N], applied as `frames @ mat.T`.
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def dft_rfft_matrix_kernel(
+    cos_ptr,      # (N_bins, N) — cos(2π k n / N)
+    nsin_ptr,     # (N_bins, N) — -sin(2π k n / N)
+    N_bins, N,
+    two_pi_over_N,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Forward rfft basis: X[k] = Σ_n x[n] (cos - i sin)(2π k n / N).
+
+    cos[k,n] = cos(2π k n / N), nsin[k,n] = -sin(2π k n / N) so that
+    X_real = x @ cos.T and X_imag = x @ nsin.T.
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    offs = 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)
+    total = N_bins * N
+    mask = offs < total
+    k = offs // N
+    n = offs % N
+    angle = (k * n).to(tl.float32) * two_pi_over_N
+    tl.store(cos_ptr + offs, tl.cos(angle), mask=mask)
+    tl.store(nsin_ptr + offs, -tl.sin(angle), mask=mask)
+
+
+@triton.jit
+def idft_c2r_matrix_kernel(
+    cos_ptr,      # (N_bins, N) — weight_k cos(2π k n / N) / N
+    sin_ptr,      # (N_bins, N) — -weight_k sin(2π k n / N) / N
+    N_bins, N,
+    two_pi_over_N,
+    inv_N,
+    n_even,       # 1 if N is even (Nyquist bin present), else 0
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Inverse rfft (irfft) basis, exploiting Hermitian symmetry.
+
+    x[n] = (1/N) Σ_k weight_k [Re(X[k]) cos(2π k n / N) - Im(X[k]) sin(2π k n / N)]
+    weight_k = 1 for k=0 and (N even) k=N/2, else 2. So
+    real_out = X_real @ cos + X_imag @ sin.
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    offs = 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)
+    total = N_bins * N
+    mask = offs < total
+    k = offs // N
+    n = offs % N
+    is_edge = (k == 0) | ((n_even == 1) & (k == N_bins - 1))
+    weight = tl.where(is_edge, 1.0, 2.0) * inv_N
+    angle = (k * n).to(tl.float32) * two_pi_over_N
+    tl.store(cos_ptr + offs, weight * tl.cos(angle), mask=mask)
+    tl.store(sin_ptr + offs, -weight * tl.sin(angle), mask=mask)