nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/moe_align.py

@@ -0,0 +1,195 @@
+"""Device-side MoE token alignment — constant-tuple, deterministic.
+
+Replaces the host moe_align (D2H of topk_ids -> numpy bincount/scatter
+-> H2D of tables) with three tiny Triton kernels so the whole routing
+pipeline stays on device. Two independent wins:
+
+1. Replay eligibility (P-REPLAY-KV-DECODE): the host align was a D2H
+   read of router output per MoE layer per step — a structural
+   plan-breaker (recorded H2D tables would replay ONE step's expert
+   choice forever; the loud refusal at recording caught exactly this).
+   Device-side, the recorded launches recompute the tables from the
+   CURRENT router output at every replay.
+2. Sync elimination: 48 D2H round-trips per decode step disappear from
+   the MoE band (same class as the device-scalar workstream).
+
+Pattern source: vLLM's moe_align_block_size CUDA kernel (grid sized on
+the worst case, consumers early-exit on the device num_tokens_post_
+padded scalar — that guard already exists in all three fused variants).
+vLLM's atomic rank assignment is NOT reproduced: atomics permute tokens
+within an expert run-to-run. Here the rank is a deterministic pairwise
+count (#earlier tokens with the same expert), which reproduces the host
+sort's order EXACTLY — the byte gate against the removed host path
+adjudicates the swap directly. n = M * top_k is tiny (8 at decode), so
+the O(n^2) rank costs nothing at the scale it runs.
+
+Outputs are sized on the shape-derived worst case
+min(n * BS, n + E * (BS - 1)) rounded up to BS — the true total lives
+ONLY in the device scalar. All launch tuples are constant per shape
+class by construction.
+"""
+
+import triton
+import triton.language as tl
+
+from nbxtensor.ops._configs import fits_shared_memory, nbx_autotune
+
+# The chunk widths and warps of the three stages, candidates the certification measures per key; the launch before
+# 2026-11-04 (BT 128, BF the experts' padded count, BLK 128 positions, BLKT = BN = 128, 4 warps — 32 and 64 wide on
+# Metal) is among them. Every candidate computes the same integers: the choice is speed only.
+_STAGE1_CONFIGS = [triton.Config({'BT': bt, 'BF': bf}, num_warps=w, num_stages=1)
+                   for bt in (32, 64, 128, 256) for bf in (16, 32, 64, 128, 256) for w in (2, 4, 8)]
+_STAGE2_CONFIGS = [triton.Config({'BLK': blk, 'BF': bf}, num_warps=w, num_stages=1)
+                   for blk in (64, 128, 256) for bf in (16, 32, 64, 128, 256) for w in (2, 4, 8)]
+_STAGE3_CONFIGS = [triton.Config({'BLKT': t, 'BN': bn}, num_warps=w, num_stages=1)
+                   for t in (32, 64, 128) for bn in (32, 64, 128) for w in (2, 4, 8)]
+
+
+def _fits(blocks) -> bool:
+    """Whether the stage's widest int32 block fits the device's shared memory per block (a cross-warp reduction stages
+    its block there; Metal, 32 KB: a [128, 128] block asked 65 536 bytes, Qwen3-30B-A3B)."""
+    return fits_shared_memory(4 * max(a * b for a, b in blocks))
+
+
+def _prune_stage1(configs, named_args, **kwargs):
+    """An expert chunk dividing the experts' padded count (both powers of two: no wider than it), blocks that fit."""
+    be = int({**named_args, **kwargs}['BE'])        # a constexpr arrives among the launch's keywords
+    return [c for c in configs if c.kwargs['BF'] <= be
+            and _fits(((c.kwargs['BF'], c.kwargs['BT']), (be, c.kwargs['BF'])))]
+
+
+def _prune_stage2(configs, named_args, **kwargs):
+    """Positions per program a whole number of blocks, an expert chunk dividing the padded count, blocks that fit."""
+    a = {**named_args, **kwargs}
+    be, bs = int(a['BE']), int(a['BS'])
+    return [c for c in configs if c.kwargs['BF'] <= be and c.kwargs['BLK'] % bs == 0
+            and _fits(((c.kwargs['BLK'], c.kwargs['BF']),))]
+
+
+def _prune_stage3(configs, named_args, **kwargs):
+    """The rank matrix [BLKT, BN] fits."""
+    return [c for c in configs if _fits(((c.kwargs['BLKT'], c.kwargs['BN']),))]
+
+
+@nbx_autotune(configs=_STAGE1_CONFIGS, key=['BS', 'BE', 'E', 'N_BUCKET'],
+              prune_configs_by={'early_config_prune': _prune_stage1}, cache_results=True)
+@triton.jit
+def moe_align_stage1_kernel(
+    topk_ids_ptr,        # [n] int64 — flat expert index per (token, k)
+    offsets_ptr,         # [E] int64 out — exclusive padded starts
+    padded_ptr,          # [E] int64 out — per-expert padded count
+    num_post_pad_ptr,    # [1] int64 out — total padded entries
+    n,
+    N_BUCKET,            # the layer's (token, expert) pairs' class (`launch_keys.moe_align_classes`): the key's (unused here)
+    BS: tl.constexpr,    # block size (fused kernel BLOCK_SIZE_M)
+    BE: tl.constexpr,    # padded expert count (pow2 >= E)
+    E: tl.constexpr,     # true expert count
+    BT: tl.constexpr,    # token chunk width
+    BF: tl.constexpr,    # expert chunk width (divides BE)
+):
+    """Single program: per-expert counts -> padded counts -> exclusive
+    prefix offsets (masked triangular sum — no tl.cumsum dependency).
+    Worked over expert chunks of BF: each chunk's padded counts go to the
+    prefix sum as a [BE, BF] block, so no [BE, BE] or [BE, BT] matrix is
+    held; BT and BF are the key's certified configuration, among the widths
+    whose blocks fit the device's shared memory (`_prune_stage1`)."""
+    e = tl.arange(0, BE)
+    excl = tl.zeros((BE,), dtype=tl.int32)
+    total = tl.sum(tl.zeros((BF,), dtype=tl.int32), axis=0)
+    for f0 in tl.static_range(0, BE, BF):
+        f = f0 + tl.arange(0, BF)
+        counts = tl.zeros((BF,), dtype=tl.int32)
+        for start in range(0, n, BT):
+            t = start + tl.arange(0, BT)
+            ids = tl.load(topk_ids_ptr + t, mask=t < n, other=-1)
+            m = (ids[None, :] == f[:, None]) & (t[None, :] < n)
+            counts += tl.sum(m.to(tl.int32), axis=1)
+        padded = ((counts + BS - 1) // BS) * BS
+        mask_f = f < E
+        tl.store(padded_ptr + f, padded.to(tl.int64), mask=mask_f)
+        total += tl.sum(tl.where(mask_f, padded, 0), axis=0)
+        tri = f[None, :] < e[:, None]
+        excl += tl.sum(tl.where(tri, padded[None, :], 0), axis=1)
+    tl.store(offsets_ptr + e, excl.to(tl.int64), mask=e < E)
+    tl.store(num_post_pad_ptr, total.to(tl.int64))
+
+
+@nbx_autotune(configs=_STAGE2_CONFIGS, key=['BS', 'BE', 'E', 'N_BUCKET'],
+              prune_configs_by={'early_config_prune': _prune_stage2}, cache_results=True)
+@triton.jit
+def moe_align_stage2_kernel(
+    offsets_ptr,         # [E] int64
+    padded_ptr,          # [E] int64
+    sorted_ids_ptr,      # [max_total] int64 out — sentinel-filled here
+    expert_ids_ptr,      # [max_blocks] int64 out — expert per block, -1 pad
+    n,
+    max_total,
+    N_BUCKET,            # the autotune key's length class (unused in the body)
+    BS: tl.constexpr,
+    BE: tl.constexpr,
+    E: tl.constexpr,
+    BLK: tl.constexpr,   # positions per program (multiple of BS)
+    BF: tl.constexpr,    # expert chunk width (divides BE)
+):
+    """Grid over output positions: sentinel-fill sorted ids and resolve
+    each block's owning expert (offsets[e] <= block_start < offsets[e]
+    + padded[e]); blocks past the true total get -1 and are skipped by
+    the fused kernels' num_tokens_post_padded early-exit anyway. The
+    experts are read in chunks of BF ([BLK, BF] blocks, the key's certified
+    configuration among those that fit, as stage 1)."""
+    pid = tl.program_id(0).to(tl.int64)
+    p = pid * BLK + tl.arange(0, BLK)
+    in_p = p < max_total
+    tl.store(sorted_ids_ptr + p, tl.zeros((BLK,), dtype=tl.int64) + n,
+             mask=in_p)
+    b = p // BS
+    pb = (b * BS).to(tl.int64)
+    eid = tl.zeros((BLK,), dtype=tl.int32)
+    hits = tl.zeros((BLK,), dtype=tl.int32)
+    for f0 in tl.static_range(0, BE, BF):
+        e = f0 + tl.arange(0, BF)
+        mask_e = e < E
+        offs = tl.load(offsets_ptr + e, mask=mask_e, other=0)
+        padd = tl.load(padded_ptr + e, mask=mask_e, other=0)
+        m = (pb[:, None] >= offs[None, :]) \
+            & (pb[:, None] < (offs + padd)[None, :]) & mask_e[None, :]
+        eid += tl.sum(tl.where(m, e[None, :], 0), axis=1)
+        hits += tl.sum(m.to(tl.int32), axis=1)
+    eid = tl.where(hits > 0, eid, -1)
+    write_b = in_p & (p % BS == 0)
+    tl.store(expert_ids_ptr + b, eid.to(tl.int64), mask=write_b)
+
+
+@nbx_autotune(configs=_STAGE3_CONFIGS, key=['E', 'N_BUCKET'],
+              prune_configs_by={'early_config_prune': _prune_stage3}, cache_results=True)
+@triton.jit
+def moe_align_stage3_kernel(
+    topk_ids_ptr,        # [n] int64
+    offsets_ptr,         # [E] int64
+    sorted_ids_ptr,      # [max_total] int64 — scatter over the sentinels
+    n,
+    N_BUCKET,            # the autotune key's length class (unused in the body)
+    E: tl.constexpr,     # true expert count (out-of-range ids skipped,
+                         # mirroring the host scatter's bounds check)
+    BLKT: tl.constexpr,  # tokens per program
+    BN: tl.constexpr,    # comparison chunk width
+):
+    """Grid over tokens: deterministic rank (#earlier tokens with the
+    same expert) -> scatter token index to offsets[expert] + rank.
+    Order-preserving, no atomics — byte-identical to the host sort it
+    replaces. MUST launch after stage 2 (stream-serialized): it
+    overwrites sentinel positions."""
+    pid = tl.program_id(0).to(tl.int64)
+    i = pid * BLKT + tl.arange(0, BLKT)
+    in_i = i < n
+    my = tl.load(topk_ids_ptr + i, mask=in_i, other=-1)
+    in_i = in_i & (my >= 0) & (my < E)
+    rank = tl.zeros((BLKT,), dtype=tl.int32)
+    for start in range(0, n, BN):
+        j = start + tl.arange(0, BN)
+        idj = tl.load(topk_ids_ptr + j, mask=j < n, other=-2)
+        m = (idj[None, :] == my[:, None]) & (j[None, :] < i[:, None])
+        rank += tl.sum(m.to(tl.int32), axis=1)
+    off_e = tl.load(offsets_ptr + my, mask=in_i, other=0)
+    pos = off_e + rank
+    tl.store(sorted_ids_ptr + pos, i.to(tl.int64), mask=in_i)

nbxtensor/ops/moe_decode_vec.py

@@ -0,0 +1,374 @@
+"""Vector (SIMT) MoE decode band — M=1, int4 g128 asym, no tl.dot.
+
+Fourth member of the SIMT decode family (decode_attn_vec + gemv_vec,
+both adopted). The decode MoE band ran as three Grouped-GEMM launches
+over sorted-token tables: at M=1 with 8 active experts of 128, every
+16-row MMA tile carries ONE real row (15/16 M-waste), plus the 3-stage
+moe_align table build, a separate SwiGLU pass, and a sum_wrapper
+combine — 18.4 ms of a 38.4 ms step for ~944 MB of weight reads
+(~12x the streaming floor; full_step_attribution_2026_08_24).
+
+The 2026-08-24 web research (R16) confirms the regime: vLLM itself
+abandoned its grouped Triton kernel at small m (its CUDA moe_wna16 —
+SIMT half2 FMA, one thread per output column, group-boundary scale
+loads — wins +41% at bs=1 on A100/DeepSeek; our num_seqs*top_k/E =
+0.0625 is two orders below its own switchover). Its combine uses
+atomicAdd (order-nondeterministic) — REFUSED here; the deterministic
+production forms fuse only the routing-weight MULTIPLY into the
+down epilogue (vLLM Triton MUL_ROUTED_WEIGHT) and keep the sum in
+fixed order. This kernel pair goes one step further and keeps the
+cross-expert sum INSIDE the down kernel as a FIXED-ORDER expert loop —
+no intermediate per-slot outputs, no sum_wrapper, no atomics.
+
+Three kernels replace the whole band (second pass ADOPTED 2026-08-25):
+  1. `moe_gateup_vec_kernel` — grid (TOP_K, cdiv(INTER, BLOCK_N)):
+     each program owns one active expert's tile of the intermediate
+     dim, runs TWO dequant-matvec accumulators (gate and up — the
+     llama.cpp dual-accumulator GLU shape, no weight-layout change)
+     over K with int4-g128-asym dequant in the FMA loop, and applies
+     SwiGLU in the epilogue: h[e, tile] = silu(g) * u.
+  2. `moe_down_split_vec_kernel` — grid (TOP_K, cdiv(HID, BLOCK_N)):
+     the expert axis lives on the GRID (8x the programs of the fused
+     loop — the occupancy lever: 32 -> 512 programs on 80 SMs moved
+     the down projection from 46 to ~248 GB/s), each program writing
+     its ROUTER-WEIGHTED partial.
+  3. `moe_part_reduce_kernel` — fixed-order sum of the TOP_K weighted
+     partials; the combine stays deterministic by construction.
+  `moe_down_combine_vec_kernel` (the first-pass fused down + combine)
+  remains as the NBX_MOEV_DSPLIT=0 kill path. moe_align disappears at
+  M=1 (no token sort exists to build); the SwiGLU pass, the
+  sum_wrapper combine, and the zero-init of the grouped caches
+  disappear with it. A gate/up-split projection variant (one matrix
+  per program + separate silu-mul epilogue) was profiled and REFUSED:
+  2x42.9+3.4 us vs the fused dual-accumulator gateup's 71 us at BN=32
+  (nsys ladder 2026-08-25). A wide-load variant (one packed int32
+  load per word, [KP, PACK, BN] register unpack) was also REFUSED:
+  x0.69 at equal blocks, best point x1.02 = rank noise — Volta's L1
+  serves the offs_k//PACK redundant requests nearly free, and the
+  cross-kernel comparison (down single-stream ~248 GB/s vs gateup
+  dual-stream ~200 at identical dequant ALU per byte) proves the
+  unpack ALU is NOT the binding constraint. DOCUMENTED WALL
+  (2026-08-25): the band's residual over its ~3.2 ms ALU working bar
+  is the dual-stream SwiGLU gateup structure at M=1 on sm_70; the
+  upstream reference for this exact regime (vLLM moe_wna16) is a CUDA
+  half2 SIMT kernel — outside Triton-pure scope. Evidence:
+  validation_outputs/moev2adopt_offbyte_gate_2026_08_25/WALL.md.
+
+Dequant expression TEXTUALLY IDENTICAL to dequant_gemv_int4_kernel
+(the parity-derived brick: q unpacked LSB-nibble from int32,
+w = q*scale + qmin in pure fp32, masked tails contribute exact +0.0).
+Expert weight/scale/qmin pointers come from the existing PtrTables
+int64 tables (tl.cast(..., tl.pointer_type, bitcast=True) — the
+established idiom).
+"""
+
+import triton
+import triton.language as tl
+
+from nbxtensor.ops._configs import fits_shared_memory, nbx_autotune
+
+# The half-precision band's tiles, candidates the certification measures per key: the ladder the band was first
+# measured over (BLOCK_N {16, 32, 64} x BLOCK_K {64, 128, 256} x warps {2, 4, 8}, Qwen3-30B-A3B's experts on a V100,
+# 2026-10-28: 193-240 us a call, the leading group around 32 / 128 / 4), and the partials' fixed-order sum's.
+_HALF_VEC_CONFIGS = [triton.Config({'BLOCK_N': bn, 'BLOCK_K': bk}, num_warps=w)
+                     for bn in (16, 32, 64) for bk in (64, 128, 256) for w in (2, 4, 8)]
+_PART_REDUCE_CONFIGS = [triton.Config({'BLOCK_N': bn}, num_warps=w) for bn in (128, 256, 512, 1024) for w in (1, 2, 4)]
+
+
+def _prune_half_vec(configs, named_args, **kwargs):
+    """The [BLOCK_K, BLOCK_N] fp32 product block each step reduces across its warps fits the device's shared memory
+    (Metal, 32 KB: 256 x 64 asked 65 536 bytes and its emitter aborted the certifier, the Mac's RUN 24d, 2026-11-04)."""
+    return [c for c in configs if fits_shared_memory(4 * c.kwargs['BLOCK_K'] * c.kwargs['BLOCK_N'])]
+
+
+@triton.jit
+def moe_gateup_vec_kernel(
+    x_ptr,            # [K] activation (fp16/fp32)
+    topk_ids_ptr,     # [TOP_K] int — active expert ids, fixed order
+    gate_qw_tab,      # [E] int64 — qweight ptrs   (int32 [K//8, N])
+    gate_sc_tab,      # [E] int64 — scales ptrs    (fp16 [G, N])
+    gate_mn_tab,      # [E] int64 — qmins ptrs     (fp16 [G, N])
+    up_qw_tab, up_sc_tab, up_mn_tab,
+    h_ptr,            # [TOP_K, N] fp32 out — silu(gate) * up
+    K, N,
+    stride_wk, stride_wn,     # packed-int32 strides (shared gate/up)
+    stride_sg, stride_sn,     # scales/qmins strides
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,    # divides GROUP (128)
+    GROUP: tl.constexpr,
+    PACK: tl.constexpr,       # 8 nibbles per int32
+):
+    pid_e = tl.program_id(0).to(tl.int64)
+    pid_n = tl.program_id(1).to(tl.int64)
+    eid = tl.load(topk_ids_ptr + pid_e).to(tl.int64)
+
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask_n = offs_n < N
+
+    g_qw = tl.cast(tl.load(gate_qw_tab + eid),
+                   tl.pointer_type(tl.int32), bitcast=True)
+    g_sc = tl.cast(tl.load(gate_sc_tab + eid),
+                   tl.pointer_type(tl.float16), bitcast=True)
+    g_mn = tl.cast(tl.load(gate_mn_tab + eid),
+                   tl.pointer_type(tl.float16), bitcast=True)
+    u_qw = tl.cast(tl.load(up_qw_tab + eid),
+                   tl.pointer_type(tl.int32), bitcast=True)
+    u_sc = tl.cast(tl.load(up_sc_tab + eid),
+                   tl.pointer_type(tl.float16), bitcast=True)
+    u_mn = tl.cast(tl.load(up_mn_tab + eid),
+                   tl.pointer_type(tl.float16), bitcast=True)
+
+    acc_g = tl.zeros((BLOCK_N,), dtype=tl.float32)
+    acc_u = tl.zeros((BLOCK_N,), dtype=tl.float32)
+    for kc in range(0, tl.cdiv(K, BLOCK_K)):
+        offs_k = kc * BLOCK_K + tl.arange(0, BLOCK_K)
+        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).to(tl.float32)
+        g = kc * BLOCK_K // GROUP  # BLOCK_K divides GROUP
+        w_off = ((offs_k[:, None] // PACK) * stride_wk
+                 + offs_n[None, :] * stride_wn)
+        shift = (offs_k[:, None] % PACK) * 4
+
+        packed = tl.load(g_qw + w_off, mask=mask, other=0)
+        q = (packed >> shift) & 0xF
+        scale = tl.load(g_sc + g * stride_sg + offs_n * stride_sn,
+                        mask=mask_n, other=0.0)
+        mn = tl.load(g_mn + g * stride_sg + offs_n * stride_sn,
+                     mask=mask_n, other=0.0)
+        w = (q.to(tl.float32) * scale[None, :].to(tl.float32)
+             + mn[None, :].to(tl.float32))
+        b = tl.where(mask, w, 0.0)
+        acc_g += tl.sum(a[:, None] * b, axis=0)
+
+        packed = tl.load(u_qw + w_off, mask=mask, other=0)
+        q = (packed >> shift) & 0xF
+        scale = tl.load(u_sc + g * stride_sg + offs_n * stride_sn,
+                        mask=mask_n, other=0.0)
+        mn = tl.load(u_mn + g * stride_sg + offs_n * stride_sn,
+                     mask=mask_n, other=0.0)
+        w = (q.to(tl.float32) * scale[None, :].to(tl.float32)
+             + mn[None, :].to(tl.float32))
+        b = tl.where(mask, w, 0.0)
+        acc_u += tl.sum(a[:, None] * b, axis=0)
+
+    # SwiGLU epilogue: silu(gate) * up, pure fp32.
+    h = acc_g * tl.sigmoid(acc_g) * acc_u
+    tl.store(h_ptr + pid_e * N + offs_n, h, mask=mask_n)
+
+
+@triton.jit
+def moe_down_combine_vec_kernel(
+    h_ptr,            # [TOP_K, Kd] fp32 — kernel 1's output
+    topk_ids_ptr,     # [TOP_K] int
+    topk_w_ptr,       # [TOP_K] routing weights (fp16/fp32)
+    down_qw_tab, down_sc_tab, down_mn_tab,
+    out_ptr,          # [N] out (fp32)
+    Kd, N,
+    stride_wk, stride_wn,
+    stride_sg, stride_sn,
+    TOP_K: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP: tl.constexpr,
+    PACK: tl.constexpr,
+):
+    """The combine IS the kernel: experts iterated in FIXED order, the
+    routing weight applied as the per-expert epilogue multiply — the
+    cross-expert sum is deterministic by construction, and no
+    intermediate per-slot output or separate reduction exists."""
+    pid_n = tl.program_id(0).to(tl.int64)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask_n = offs_n < N
+
+    acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
+    for e in tl.static_range(TOP_K):
+        eid = tl.load(topk_ids_ptr + e).to(tl.int64)
+        we = tl.load(topk_w_ptr + e).to(tl.float32)
+        qw = tl.cast(tl.load(down_qw_tab + eid),
+                     tl.pointer_type(tl.int32), bitcast=True)
+        sc = tl.cast(tl.load(down_sc_tab + eid),
+                     tl.pointer_type(tl.float16), bitcast=True)
+        mp = tl.cast(tl.load(down_mn_tab + eid),
+                     tl.pointer_type(tl.float16), bitcast=True)
+        part = tl.zeros((BLOCK_N,), dtype=tl.float32)
+        for kc in range(0, tl.cdiv(Kd, BLOCK_K)):
+            offs_k = kc * BLOCK_K + tl.arange(0, BLOCK_K)
+            mask_k = offs_k < Kd
+            mask = mask_k[:, None] & mask_n[None, :]
+            a = tl.load(h_ptr + e * Kd + offs_k, mask=mask_k,
+                        other=0.0)
+            g = kc * BLOCK_K // GROUP
+            packed = tl.load(qw + (offs_k[:, None] // PACK) * stride_wk
+                             + offs_n[None, :] * stride_wn,
+                             mask=mask, other=0)
+            q = (packed >> ((offs_k[:, None] % PACK) * 4)) & 0xF
+            scale = tl.load(sc + g * stride_sg + offs_n * stride_sn,
+                            mask=mask_n, other=0.0)
+            mn = tl.load(mp + g * stride_sg + offs_n * stride_sn,
+                         mask=mask_n, other=0.0)
+            w = (q.to(tl.float32) * scale[None, :].to(tl.float32)
+                 + mn[None, :].to(tl.float32))
+            b = tl.where(mask, w, 0.0)
+            part += tl.sum(a[:, None] * b, axis=0)
+        acc += we * part
+
+    tl.store(out_ptr + offs_n, acc, mask=mask_n)
+
+
+@triton.jit
+def moe_down_split_vec_kernel(
+    h_ptr, topk_ids_ptr, topk_w_ptr,
+    down_qw_tab, down_sc_tab, down_mn_tab,
+    part_ptr,         # [TOP_K, N] fp32 — weighted per-expert partials
+    Kd, N,
+    stride_wk, stride_wn,
+    stride_sg, stride_sn,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP: tl.constexpr,
+    PACK: tl.constexpr,
+):
+    """ADOPTED second pass (2026-08-25): the expert axis moves to the
+    GRID (grid = (TOP_K, tiles) — 8x the programs of the fused loop),
+    each program writing its ROUTER-WEIGHTED partial; a separate tiny
+    fixed-order reduce completes the deterministic combine. Trades one
+    launch + a [TOP_K, N] fp32 buffer for occupancy and a straight-line
+    inner loop."""
+    pid_e = tl.program_id(0).to(tl.int64)
+    pid_n = tl.program_id(1).to(tl.int64)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask_n = offs_n < N
+    eid = tl.load(topk_ids_ptr + pid_e).to(tl.int64)
+    we = tl.load(topk_w_ptr + pid_e).to(tl.float32)
+    qw = tl.cast(tl.load(down_qw_tab + eid),
+                 tl.pointer_type(tl.int32), bitcast=True)
+    sc = tl.cast(tl.load(down_sc_tab + eid),
+                 tl.pointer_type(tl.float16), bitcast=True)
+    mp = tl.cast(tl.load(down_mn_tab + eid),
+                 tl.pointer_type(tl.float16), bitcast=True)
+    acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
+    for kc in range(0, tl.cdiv(Kd, BLOCK_K)):
+        offs_k = kc * BLOCK_K + tl.arange(0, BLOCK_K)
+        mask_k = offs_k < Kd
+        mask = mask_k[:, None] & mask_n[None, :]
+        a = tl.load(h_ptr + pid_e * Kd + offs_k, mask=mask_k, other=0.0)
+        g = kc * BLOCK_K // GROUP
+        packed = tl.load(qw + (offs_k[:, None] // PACK) * stride_wk
+                         + offs_n[None, :] * stride_wn, mask=mask, other=0)
+        q = (packed >> ((offs_k[:, None] % PACK) * 4)) & 0xF
+        scale = tl.load(sc + g * stride_sg + offs_n * stride_sn,
+                        mask=mask_n, other=0.0)
+        mn = tl.load(mp + g * stride_sg + offs_n * stride_sn,
+                     mask=mask_n, other=0.0)
+        w = (q.to(tl.float32) * scale[None, :].to(tl.float32)
+             + mn[None, :].to(tl.float32))
+        b = tl.where(mask, w, 0.0)
+        acc += tl.sum(a[:, None] * b, axis=0)
+    tl.store(part_ptr + pid_e * N + offs_n, we * acc, mask=mask_n)
+
+
+@nbx_autotune(configs=_PART_REDUCE_CONFIGS, key=['N', 'TOP_K'], cache_results=True)
+@triton.jit
+def moe_part_reduce_kernel(
+    part_ptr, out_ptr, N,
+    TOP_K: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """Fixed-order sum of the TOP_K weighted partials — deterministic
+    (the same contract as flash_decode_reduce: order is a constant)."""
+    pid = tl.program_id(0).to(tl.int64)
+    offs_n = pid * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask_n = offs_n < N
+    acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
+    for e in tl.static_range(TOP_K):
+        acc += tl.load(part_ptr + e * N + offs_n, mask=mask_n, other=0.0)
+    tl.store(out_ptr + offs_n, acc, mask=mask_n)
+
+
+# ---------------------------------------------------------------------------
+# Half-precision experts (fp16 / bf16), the fifth member (S3, 2026-10-28).
+#
+# The same band for experts stored in half precision: at M=1 the grouped GEMM
+# ran Qwen3-30B-A3B's decode (fp16 experts, 8 of 128, 2048 x 768) at 0.35 ms a
+# call for ~24 MB of weights (~70 GB/s, nsys on 4 V100, 52 ms of GPU a token);
+# every 16-row tile carried one real row. Same grids as the int4 member — the
+# expert axis on the grid for both projections, SwiGLU in the gateup epilogue,
+# the router-weighted partials summed in fixed order by moe_part_reduce_kernel —
+# with the weight tile read as stored (no dequant) and every product accumulated
+# in fp32. W_BF16 selects the element type the pointer tables point at.
+# ---------------------------------------------------------------------------
+
+
+@nbx_autotune(configs=_HALF_VEC_CONFIGS, key=['K', 'N', 'TOP_K', 'W_BF16'],
+              prune_configs_by={'early_config_prune': _prune_half_vec}, cache_results=True)
+@triton.jit
+def moe_gateup_vec_half_kernel(
+    x_ptr,            # [K] activation
+    topk_ids_ptr,     # [TOP_K] int — active expert ids, fixed order
+    gate_tab, up_tab,  # [E] int64 — expert weight ptrs ([N, K] or [K, N] by the strides)
+    h_ptr,            # [TOP_K, N] fp32 out — silu(gate) * up
+    K, N,
+    stride_wk, stride_wn,
+    TOP_K,            # the routed experts (the grid's first axis): the autotune key's, unused in the body
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    W_BF16: tl.constexpr,
+):
+    pid_e = tl.program_id(0).to(tl.int64)
+    pid_n = tl.program_id(1).to(tl.int64)
+    eid = tl.load(topk_ids_ptr + pid_e).to(tl.int64)
+    wty = tl.bfloat16 if W_BF16 else tl.float16
+    g_w = tl.cast(tl.load(gate_tab + eid), tl.pointer_type(wty), bitcast=True)
+    u_w = tl.cast(tl.load(up_tab + eid), tl.pointer_type(wty), bitcast=True)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask_n = offs_n < N
+    acc_g = tl.zeros((BLOCK_N,), dtype=tl.float32)
+    acc_u = tl.zeros((BLOCK_N,), dtype=tl.float32)
+    for kc in range(0, tl.cdiv(K, BLOCK_K)):
+        offs_k = kc * BLOCK_K + tl.arange(0, BLOCK_K)
+        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).to(tl.float32)
+        w_off = offs_k[:, None] * stride_wk + offs_n[None, :] * stride_wn
+        g = tl.load(g_w + w_off, mask=mask, other=0.0).to(tl.float32)
+        acc_g += tl.sum(a[:, None] * g, axis=0)
+        u = tl.load(u_w + w_off, mask=mask, other=0.0).to(tl.float32)
+        acc_u += tl.sum(a[:, None] * u, axis=0)
+    h = acc_g * tl.sigmoid(acc_g) * acc_u
+    tl.store(h_ptr + pid_e * N + offs_n, h, mask=mask_n)
+
+
+@nbx_autotune(configs=_HALF_VEC_CONFIGS, key=['Kd', 'N', 'TOP_K', 'W_BF16'],
+              prune_configs_by={'early_config_prune': _prune_half_vec}, cache_results=True)
+@triton.jit
+def moe_down_split_vec_half_kernel(
+    h_ptr,            # [TOP_K, Kd] fp32 — the gateup output
+    topk_ids_ptr, topk_w_ptr,
+    down_tab,         # [E] int64 — expert down weight ptrs
+    part_ptr,         # [TOP_K, N] fp32 — router-weighted per-expert partials
+    Kd, N,
+    stride_wk, stride_wn,
+    TOP_K,            # the routed experts (the grid's first axis): the autotune key's, unused in the body
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    W_BF16: tl.constexpr,
+):
+    pid_e = tl.program_id(0).to(tl.int64)
+    pid_n = tl.program_id(1).to(tl.int64)
+    eid = tl.load(topk_ids_ptr + pid_e).to(tl.int64)
+    rw = tl.load(topk_w_ptr + pid_e).to(tl.float32)
+    wty = tl.bfloat16 if W_BF16 else tl.float16
+    d_w = tl.cast(tl.load(down_tab + eid), tl.pointer_type(wty), bitcast=True)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask_n = offs_n < N
+    acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
+    for kc in range(0, tl.cdiv(Kd, BLOCK_K)):
+        offs_k = kc * BLOCK_K + tl.arange(0, BLOCK_K)
+        mask_k = offs_k < Kd
+        mask = mask_k[:, None] & mask_n[None, :]
+        a = tl.load(h_ptr + pid_e * Kd + offs_k, mask=mask_k, other=0.0)
+        w_off = offs_k[:, None] * stride_wk + offs_n[None, :] * stride_wn
+        d = tl.load(d_w + w_off, mask=mask, other=0.0).to(tl.float32)
+        acc += tl.sum(a[:, None] * d, axis=0)
+    tl.store(part_ptr + pid_e * N + offs_n, acc * rw, mask=mask_n)