From 6bf1509c1ac9337aaea9a1c776851e13da21aa09 Mon Sep 17 00:00:00 2001 From: XiaobingSuper Date: Fri, 31 Jul 2026 11:17:22 +0000 Subject: [PATCH 1/5] gfx1250: scan every expert in the grouped-MoE contiguous-M psum The tile-aligned prefix scan ran one thread per expert in a single MAX_EXPERTS_PER_BLOCK (512) block, so any model with more experts than that silently lost the tail. Kimi-K3 has 896: experts 512..895 were never scanned, `starts`/`psum`/`contiguous_m` kept the uninitialised values of their torch.empty allocation, the masked-to-contiguous row remap turned those into out-of-range rows, and the downstream moe_fused_quant_preshuffle_routeks_* faulted on them (HSA_STATUS_ERROR_MEMORY_FAULT). The cap was known -- the fused route+psum variant is gated on _FUSED_ROUTE_PSUM_MAX_EXPERTS -- but this path had no such guard and corrupted quietly instead. Sweep the experts in block-sized chunks and carry the running offset between them in LDS, so E is no longer bounded by the block width. The carry has to live in LDS rather than a register because the chunk loop is a runtime loop. Lanes past `experts` now feed 0 into the scan so the last lane still holds the chunk total, and the two kernels spell the sweep out separately because @flyc.kernel only AST-transforms the decorated body. Checked against a torch.cumsum reference over the tile-aligned counts for E = 8/256/512/896/1024: exact, where before E=896 mismatched on exactly 384 experts starting at 512. --- .../ops/flydsl/kernels/moe_contiguous_psum.py | 162 ++++++++++++------ 1 file changed, 106 insertions(+), 56 deletions(-) diff --git a/aiter/ops/flydsl/kernels/moe_contiguous_psum.py b/aiter/ops/flydsl/kernels/moe_contiguous_psum.py index ca84b12ffa..e6b0493c9a 100644 --- a/aiter/ops/flydsl/kernels/moe_contiguous_psum.py +++ b/aiter/ops/flydsl/kernels/moe_contiguous_psum.py @@ -6,6 +6,10 @@ Computes tile-aligned exclusive prefix sum of per-expert counts for the contiguous grouped-GEMM scheduler. Single-block parallel scan replaces torch.cumsum (avoids rocprim trampoline overhead for small E). + +The block is ``MAX_EXPERTS_PER_BLOCK`` threads wide but E is not bounded by it: +the scan sweeps the experts in block-sized chunks and carries the running offset +between chunks in LDS. Kimi-K3 (E=896) is the first model to exceed one chunk. """ import flydsl.compiler as flyc @@ -29,12 +33,14 @@ class _PsumStorage: """LDS for the prefix-scan kernels. ``lds0``/``lds1`` are ping-pong buffers: each Hillis-Steele step reads one - and writes the other, then the two swap. The trailing 16 is the byte - alignment of each array. + and writes the other, then the two swap. ``carry`` accumulates the total of + the chunks already scanned, so an E wider than the block still gets one + continuous prefix sum. The trailing 16 is the byte alignment of each array. """ lds0: fx.Array[fx.Int32, MAX_EXPERTS_PER_BLOCK, 16] lds1: fx.Array[fx.Int32, MAX_EXPERTS_PER_BLOCK, 16] + carry: fx.Array[fx.Int32, 1, 16] @fx.struct @@ -60,6 +66,11 @@ def _lds_store(ptr, val, idx): fx.ptr_store(val, ptr + fx.Int64(idx)) +# The chunked scan below is written out in both kernels rather than shared: +# @flyc.kernel AST-transforms only the decorated body, so a dynamic `for`/`if` +# does not survive being factored into a plain helper. + + def build_moe_contiguous_psum_module(): """JIT launcher: tile-aligned prefix sum over per-expert counts.""" @@ -86,50 +97,69 @@ def psum_kernel( lds = fx.SharedAllocator().allocate(_PsumStorage).peek() lds0 = lds.lds0.ptr lds1 = lds.lds1.ptr + carry = lds.carry.ptr m_rsrc = ptr_rsrc(masked_m) s_rsrc = ptr_rsrc(starts) p_rsrc = ptr_rsrc(psum) c_rsrc = ptr_rsrc(contiguous_m) - in_range = tid < fx.Uint32(experts) - if in_range: - m = fx.Uint32(buffer_ops.buffer_load(m_rsrc, tid, vec_width=1, dtype=i32)) - _lds_store(lds0, (m + tile_minus_1) // tile_v * tile_v, tid) - + is_lane0 = tid == fx.Uint32(0) + if is_lane0: + _lds_store(carry, fx.Int32(0), 0) gpu.barrier() - src = lds0 - dst = lds1 - for offset in range_constexpr(1, MAX_EXPERTS_PER_BLOCK): - if const_expr((offset & (offset - 1)) != 0): - continue - if in_range: + # E is not bounded by the block width: sweep the experts in block-sized + # chunks, offsetting each chunk by the running total of the ones before + # it. That total lives in LDS (``carry``) because a register would not + # survive the runtime chunk loop. Lanes past ``experts`` feed 0 into the + # scan so the last lane still holds the chunk total, and write nothing. + for base in range(0, experts, MAX_EXPERTS_PER_BLOCK): + e = fx.Uint32(base) + tid + in_expert = e < fx.Uint32(experts) + m_e = fx.Uint32(0) + if in_expert: + m_e = fx.Uint32( + buffer_ops.buffer_load(m_rsrc, e, vec_width=1, dtype=i32) + ) + _lds_store(lds0, fx.Int32((m_e + tile_minus_1) // tile_v * tile_v), tid) + gpu.barrier() + + src = lds0 + dst = lds1 + for offset in range_constexpr(1, MAX_EXPERTS_PER_BLOCK): + if const_expr((offset & (offset - 1)) != 0): + continue val = _lds_load(src, tid) has_prev = tid >= offset prev = fx.Int32(0) if has_prev: prev = _lds_load(src, tid - offset) _lds_store(dst, val + prev, tid) - gpu.barrier() - src, dst = dst, src + gpu.barrier() + src, dst = dst, src - if in_range: - is_not_first = tid != 0 - start = fx.Int32(0) - if is_not_first: - start = _lds_load(src, tid - 1) - m_tid = fx.Int32( - buffer_ops.buffer_load(m_rsrc, tid, vec_width=1, dtype=i32) - ) - buffer_ops.buffer_store(start, s_rsrc, tid) - buffer_ops.buffer_store(start + m_tid, p_rsrc, tid) + base_off = _lds_load(carry, 0) + if in_expert: + is_not_first = tid != 0 + excl = fx.Int32(0) + if is_not_first: + excl = _lds_load(src, tid - 1) + start = excl + base_off + buffer_ops.buffer_store(start, s_rsrc, e) + buffer_ops.buffer_store(start + fx.Int32(m_e), p_rsrc, e) + + # Fold this chunk's total in before the next one overwrites lds0. + chunk_total = _lds_load(src, MAX_EXPERTS_PER_BLOCK - 1) + gpu.barrier() + if is_lane0: + _lds_store(carry, base_off + chunk_total, 0) + gpu.barrier() - is_last = tid == fx.Uint32(experts) - 1 - if is_last: - final_cur = _lds_load(src, tid) - gt = final_cur > fx.Int32(tile_v) - buffer_ops.buffer_store(gt.select(final_cur, tile_v), c_rsrc, 0) + if is_lane0: + total = _lds_load(carry, 0) + gt = total > fx.Int32(tile_v) + buffer_ops.buffer_store(gt.select(total, tile_v), c_rsrc, 0) @flyc.jit def launch_psum( @@ -187,6 +217,7 @@ def psum_remap_kernel( lds = fx.SharedAllocator().allocate(_PsumStorage).peek() lds0 = lds.lds0.ptr lds1 = lds.lds1.ptr + carry = lds.carry.ptr m_rsrc = ptr_rsrc(masked_m) rows_rsrc = ptr_rsrc(topids_to_rows) @@ -194,43 +225,62 @@ def psum_remap_kernel( p_rsrc = ptr_rsrc(psum) c_rsrc = ptr_rsrc(contiguous_m) - in_expert = tid < fx.Uint32(experts) - if in_expert: - m = fx.Uint32(buffer_ops.buffer_load(m_rsrc, tid, vec_width=1, dtype=i32)) - _lds_store(lds0, (m + tile_minus_1) // tile_v * tile_v, tid) - + is_lane0 = tid == fx.Uint32(0) + if is_lane0: + _lds_store(carry, fx.Int32(0), 0) gpu.barrier() - src = lds0 - dst = lds1 - for offset in range_constexpr(1, MAX_EXPERTS_PER_BLOCK): - if const_expr((offset & (offset - 1)) != 0): - continue + # E is not bounded by the block width: sweep the experts in block-sized + # chunks, offsetting each chunk by the running total of the ones before + # it. That total lives in LDS (``carry``) because a register would not + # survive the runtime chunk loop. Lanes past ``experts`` feed 0 into the + # scan so the last lane still holds the chunk total, and write nothing. + for base in range(0, experts, MAX_EXPERTS_PER_BLOCK): + e = fx.Uint32(base) + tid + in_expert = e < fx.Uint32(experts) + m_e = fx.Uint32(0) if in_expert: + m_e = fx.Uint32( + buffer_ops.buffer_load(m_rsrc, e, vec_width=1, dtype=i32) + ) + _lds_store(lds0, fx.Int32((m_e + tile_minus_1) // tile_v * tile_v), tid) + gpu.barrier() + + src = lds0 + dst = lds1 + for offset in range_constexpr(1, MAX_EXPERTS_PER_BLOCK): + if const_expr((offset & (offset - 1)) != 0): + continue val = _lds_load(src, tid) has_prev = tid >= offset prev = fx.Int32(0) if has_prev: prev = _lds_load(src, tid - offset) _lds_store(dst, val + prev, tid) + gpu.barrier() + src, dst = dst, src + + base_off = _lds_load(carry, 0) + if in_expert: + is_not_first = tid != 0 + excl = fx.Int32(0) + if is_not_first: + excl = _lds_load(src, tid - 1) + start = excl + base_off + buffer_ops.buffer_store(start, s_rsrc, e) + buffer_ops.buffer_store(start + fx.Int32(m_e), p_rsrc, e) + + # Fold this chunk's total in before the next one overwrites lds0. + chunk_total = _lds_load(src, MAX_EXPERTS_PER_BLOCK - 1) + gpu.barrier() + if is_lane0: + _lds_store(carry, base_off + chunk_total, 0) gpu.barrier() - src, dst = dst, src - if in_expert: - is_not_first = tid != 0 - start = fx.Int32(0) - if is_not_first: - start = _lds_load(src, tid - 1) - m_tid = fx.Int32( - buffer_ops.buffer_load(m_rsrc, tid, vec_width=1, dtype=i32) - ) - buffer_ops.buffer_store(start, s_rsrc, tid) - buffer_ops.buffer_store(start + m_tid, p_rsrc, tid) - is_last = tid == fx.Uint32(experts) - 1 - if is_last: - final_cur = _lds_load(src, tid) - gt = final_cur > fx.Int32(tile_v) - buffer_ops.buffer_store(gt.select(final_cur, tile_v), c_rsrc, 0) + if is_lane0: + total = _lds_load(carry, 0) + gt = total > fx.Int32(tile_v) + buffer_ops.buffer_store(gt.select(total, tile_v), c_rsrc, 0) gpu.barrier() From 6c24d14c28812576f5123ad7df457143bdefbdc3 Mon Sep 17 00:00:00 2001 From: XiaobingSuper Date: Fri, 31 Jul 2026 11:17:35 +0000 Subject: [PATCH 2/5] gfx1250: compute SiTUv2 in the grouped TDM MoE stage1 epilogue The TDM stage1 act code only encoded silu and swiglu, so ActivationType.Situv2 fell through to `stage1_act = 1` and was computed as silu -- quietly, since situ_beta/situ_linear_beta were accepted at the grouped entry point and then dropped. Kimi-K3 is `hidden_act="situ"`, so on gfx1250 (where the separated path has no working SiTUv2 kernel) the model had no correct MoE at all. This is the TODO(situv2) left behind when the fused stage1 epilogue was removed. Add stage1_act=3 and wire it through both epilogues: the batched one used by the a8w4 fused-quant path and the element-wise one used by a4w4's bf16 intermediate. beta/linear_beta are runtime kernel arguments, so every SiTUv2 shape shares one compiled kernel; their reciprocals are taken on the host so folding them into the per-element multipliers stays exact, and the multipliers themselves are hoisted out of the inner loop. tanh uses the saturating identity 2*sigmoid(2z)-1 rather than the (1-e)/(1+e) form: exp2 of a large positive argument goes to +inf and rcp(+inf) to 0, so both tails are correct without an |x| fixup or a sign select. SiTUv2 is bounded by construction and takes no swiglu clamp. Un-skip the grouped SiTUv2 test and add the a8w4 case (a separate code path from a4w4); both land at rel_l2 ~3e-3 against the fp32 reference, which is MXFP4 quantisation noise, with silu and swiglu unchanged. Kimi-K3 end to end on 4xMI450 is GSM8K 1319 = 0.9591. --- aiter/ops/flydsl/batched_gemm_mxfp4.py | 26 +++- aiter/ops/flydsl/grouped_moe_gfx1250.py | 24 +++- .../ops/flydsl/kernels/gemm_common_gfx1250.py | 121 ++++++++++++++++++ .../kernels/mxfp4_preshuffle_gfx1250_tdm.py | 67 ++++++++-- aiter/ops/flydsl/moe_common.py | 5 +- op_tests/test_flydsl_grouped_gemm_gfx1250.py | 49 +++++-- 6 files changed, 256 insertions(+), 36 deletions(-) diff --git a/aiter/ops/flydsl/batched_gemm_mxfp4.py b/aiter/ops/flydsl/batched_gemm_mxfp4.py index 3f546349bb..0acf83f7fd 100644 --- a/aiter/ops/flydsl/batched_gemm_mxfp4.py +++ b/aiter/ops/flydsl/batched_gemm_mxfp4.py @@ -51,6 +51,8 @@ def flydsl_grouped_gemm_a8w4_masked( stage1_quant_out=0, quant_scale=None, quant_wmma_rep=1, + situ_beta=1.0, + situ_linear_beta=1.0, ): """Contiguous-M grouped a8w4 GEMM on the batched TDM kernel. @@ -65,16 +67,26 @@ def flydsl_grouped_gemm_a8w4_masked( m_tile_map (n_experts,) int32 psum (per-expert exclusive end-row) contiguous_m must be a multiple of tile_m (holds by construction). - When ``stage1_quant_out=1`` (fp8), the epilogue fuses silu/swiglu + MX fp8 - quantization + e8m0 scale preshuffle into the kernel. ``out`` receives the - fp8 payload (uint8, 1 byte/elem) and ``quant_scale`` receives the preshuffled - e8m0 scale (uint8). ``quant_wmma_rep`` is gemm2's ``warp_tile_m // 16``, - controlling the scale preshuffle tile geometry. + ``stage1_act`` selects the stage1 epilogue: 0 none, 1 silu, 2 swiglu, + 3 SiTUv2 (``situ_beta`` / ``situ_linear_beta``, the Kimi-K3 activation). + The betas are runtime kernel arguments, so all SiTUv2 shapes share one + compiled kernel; their reciprocals are taken here rather than by an + in-kernel v_rcp_f32 so the fold stays exact. + + When ``stage1_quant_out=1`` (fp8), the epilogue fuses the activation + MX + fp8 quantization + e8m0 scale preshuffle into the kernel. ``out`` receives + the fp8 payload (uint8, 1 byte/elem) and ``quant_scale`` receives the + preshuffled e8m0 scale (uint8). ``quant_wmma_rep`` is gemm2's + ``warp_tile_m // 16``, controlling the scale preshuffle tile geometry. """ from .kernels.mxfp4_preshuffle_gfx1250_tdm import launch_gemm_a8w4_tdm if stream is None: stream = torch.cuda.current_stream() + if float(situ_beta) <= 0.0: + raise ValueError(f"situ_beta must be > 0, got {situ_beta!r}") + if float(situ_linear_beta) <= 0.0: + raise ValueError(f"situ_linear_beta must be > 0, got {situ_linear_beta!r}") nb = min(num_buffers, max(1, K // tile_k)) has_bias = 1 if bias is not None else 0 bias_ptr = ptr_arg(bias) if bias is not None else ptr_arg(a) @@ -110,6 +122,10 @@ def flydsl_grouped_gemm_a8w4_masked( stage1_quant_out, quant_wmma_rep, quant_scale_tensor, + float(situ_beta), + 1.0 / float(situ_beta), + float(situ_linear_beta), + 1.0 / float(situ_linear_beta), ) return out diff --git a/aiter/ops/flydsl/grouped_moe_gfx1250.py b/aiter/ops/flydsl/grouped_moe_gfx1250.py index 866618d565..722f151bb2 100644 --- a/aiter/ops/flydsl/grouped_moe_gfx1250.py +++ b/aiter/ops/flydsl/grouped_moe_gfx1250.py @@ -398,6 +398,8 @@ def _grouped_a8w4_tdm_moe( data_format="a8w4", expert_mask=None, num_local_tokens=None, + situ_beta=1.0, + situ_linear_beta=1.0, ): import functools @@ -482,12 +484,22 @@ def _grouped_a8w4_tdm_moe( out_is_f16 = 1 if (dtype == torch.float16 or dtype == dtypes.fp16) else 0 two_inter = 2 * inter_dim - stage1_act = 2 if activation == ActivationType.Swiglu else 1 + # Stage1 epilogue code: 1 silu, 2 swiglu, 3 SiTUv2. The caller has already + # rejected anything else. + if activation == ActivationType.Swiglu: + stage1_act = 2 + elif activation == ActivationType.Situv2: + stage1_act = 3 + else: + stage1_act = 1 + # SiTUv2 is bounded by construction and takes no clamp, so the limit only + # ever applies to swiglu. sl = ( float(swiglu_limit) if swiglu_limit else (7.0 if activation == ActivationType.Swiglu else float("inf")) ) + _situ_kw = {"situ_beta": situ_beta, "situ_linear_beta": situ_linear_beta} _b1 = ( bias1.to(dtype).contiguous() if (bias1 is not None and bias1.numel() > 0) @@ -560,6 +572,7 @@ def _grouped_a8w4_tdm_moe( stage1_quant_out=1, quant_scale=a2_scale, quant_wmma_rep=wmma_rep2, + **_situ_kw, ) else: # Original path: bf16 intermediate + separate quant kernel. @@ -584,6 +597,7 @@ def _grouped_a8w4_tdm_moe( bias=_b1, swiglu_limit=sl, num_buffers=num_buffers, + **_situ_kw, ) a2_payload, a2_scale = flydsl_moe_fused_quant_preshuffle( y, @@ -648,6 +662,7 @@ def _grouped_a8w4_tdm_moe( stage1_quant_out=1, quant_scale=a2_scale, quant_wmma_rep=wmma_rep2, + **_situ_kw, ), ) ) @@ -676,6 +691,7 @@ def _grouped_a8w4_tdm_moe( bias=_b1, swiglu_limit=sl, num_buffers=num_buffers, + **_situ_kw, ), ) ) @@ -998,6 +1014,8 @@ def _tdm_env(name): data_format=data_format, expert_mask=expert_mask, num_local_tokens=num_local_tokens, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, **_tdm_kw, ) @@ -1005,10 +1023,6 @@ def _tdm_env(name): # GEMM (gemm_mxscale_gfx1250 / moe_grouped_gemm_mxscale_gfx1250) was # removed. Anything the TDM path cannot serve falls back to the caller's # generic MoE via None. - # TODO(situv2): ActivationType.Situv2 used to be handled by the deleted - # fused stage1 epilogue. The TDM stage1 act code only encodes silu/swiglu, - # so Situv2 currently runs as silu here -- add a real situv2 code (and - # plumb situ_beta / situ_linear_beta) when the TDM refactor settles. # TODO(aot): AOT has no coverage for the TDM batched GEMM # (batched_gemm_mxfp4); grouped kernels are JIT-compiled at first use # until that is added back. diff --git a/aiter/ops/flydsl/kernels/gemm_common_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_common_gfx1250.py index 24f0bb01b5..2343380dff 100644 --- a/aiter/ops/flydsl/kernels/gemm_common_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_common_gfx1250.py @@ -1,5 +1,7 @@ """Shared utilities for gfx1250 GEMM kernels (fp16 / mxfp4 / mxfp8).""" +from collections import namedtuple + import flydsl.expr as fx from flydsl.expr import arith, gpu, rocdl, tdm_ops from flydsl.expr.arith import _to_raw as _raw @@ -97,6 +99,121 @@ def fused_silu_swiglu_elem(g, u, *, swiglu, limit_f32, neg_limit_f32): return g * sig * u +def _tanh_f32(x, tanh_mul): + """tanh(x) via the sigmoid identity tanh(z) = 2*sigmoid(2z) - 1. + + ``tanh_mul`` is the caller-hoisted ``-2*log2(e)/beta`` multiplier, so this + evaluates ``2*rcp(1 + exp2(x * tanh_mul)) - 1`` for ``tanh(x/beta)`` in one + exp2 + one rcp. Saturating rather than branchy: a large positive argument + drives exp2 to +inf and rcp(+inf) to 0 (-> -1), a large negative one drives + exp2 to 0 (-> +1), so no |x| fixup or sign select is needed. + """ + import flydsl.expr as _fx + + _one = _fx.Float32(1.0) + _two = _fx.Float32(2.0) + exp_val = _fx.Float32(rocdl.exp2(T.f32, _raw(x * tanh_mul))) + rcp_val = _fx.Float32(rocdl.rcp(T.f32, _one + exp_val)) + return _two * rcp_val - _one + + +# Loop-invariant f32 multipliers for the SiTUv2 epilogue, hoisted out of the +# per-element math by situv2_consts(). +SituV2Consts = namedtuple("SituV2Consts", "beta gate_tanh_mul linear_beta up_tanh_mul") + + +def situv2_consts(beta, inv_beta, linear_beta, inv_linear_beta): + """Fold the SiTUv2 betas into the per-element multipliers, once per kernel. + + ``inv_beta`` / ``inv_linear_beta`` are the host-computed reciprocals (exact, + unlike an in-kernel v_rcp_f32). Everything here is uniform across the tile, + so hoisting it keeps the inner loop at 3 exp2 + 3 rcp per element. + """ + import flydsl.expr as _fx + + neg_two_log2e = _fx.Float32(-2.0 * LOG2E) + return SituV2Consts( + beta=beta, + gate_tanh_mul=neg_two_log2e * inv_beta, + linear_beta=linear_beta, + up_tanh_mul=neg_two_log2e * inv_linear_beta, + ) + + +def fused_situv2_elem(g, u, *, consts): + """One (gate, up) pair -> SiTUv2 (Kimi-K3 hidden_act="situ"). + + beta * tanh(g/beta) * sigmoid(g) * linear_beta * tanh(u/linear_beta) + + ``consts`` comes from situv2_consts(). No clamp: SiTUv2 is bounded by + construction, so the swiglu limit does not apply. + """ + import flydsl.expr as _fx + + _one = _fx.Float32(1.0) + nlog2e = _fx.Float32(-LOG2E) + exp_val = _fx.Float32(rocdl.exp2(T.f32, _raw(g * nlog2e))) + sig = _fx.Float32(rocdl.rcp(T.f32, _one + exp_val)) + gate_act = consts.beta * _tanh_f32(g, consts.gate_tanh_mul) * sig + up_act = consts.linear_beta * _tanh_f32(u, consts.up_tanh_mul) + return gate_act * up_act + + +def batched_situv2(pairs, *, consts, range_constexpr): + """Batched SiTUv2 with pipelined exp2/rcp for better TRANS utilisation. + + Same staging idea as batched_silu_swiglu, over the three transcendental + pairs SiTUv2 needs per element: sigmoid(gate), tanh(gate/beta) and + tanh(up/linear_beta). Grouping all exp2s, then all rcps, keeps the TRANS + unit busy instead of stalling on each dependent pair in turn. + + Args: + pairs: list of (gate, up) f32 value pairs. + consts: SituV2Consts from situv2_consts(). + range_constexpr: the FlyDSL ``range_constexpr`` helper. + + Returns: + list of activated f32 values, same length as *pairs*. + """ + import flydsl.expr as _fx + + _one = _fx.Float32(1.0) + _two = _fx.Float32(2.0) + nlog2e = _fx.Float32(-LOG2E) + N = len(pairs) + # Stage 1: all exp2 arguments, then all exp2. + args = [] + for i in range_constexpr(N): + g, u = pairs[i] + args.append(g * nlog2e) # sigmoid(gate) + args.append(g * consts.gate_tanh_mul) # tanh(gate/beta) + args.append(u * consts.up_tanh_mul) # tanh(up/linear_beta) + rocdl.sched_barrier(0) + exp_vals = [] + for i in range_constexpr(3 * N): + exp_vals.append(_fx.Float32(rocdl.exp2(T.f32, _raw(args[i])))) + # Stage 2a: 1 + exp + rocdl.sched_barrier(0) + sum_vals = [] + for i in range_constexpr(3 * N): + sum_vals.append(_one + exp_vals[i]) + # Stage 2b: rcp + rocdl.sched_barrier(0) + rcp_vals = [] + for i in range_constexpr(3 * N): + rcp_vals.append(_fx.Float32(rocdl.rcp(T.f32, sum_vals[i]))) + # Stage 3: sigmoid / tanh assembly and the final product. + rocdl.sched_barrier(0) + results = [] + for i in range_constexpr(N): + sig = rcp_vals[3 * i] + gate_tanh = _two * rcp_vals[3 * i + 1] - _one + up_tanh = _two * rcp_vals[3 * i + 2] - _one + gate_act = consts.beta * gate_tanh * sig + results.append(gate_act * (consts.linear_beta * up_tanh)) + return results + + def batched_silu_swiglu(pairs, *, swiglu, limit_f32, neg_limit_f32, range_constexpr): """Batched silu/swiglu with pipelined exp2/rcp for better TRANS utilisation. @@ -148,11 +265,15 @@ def batched_silu_swiglu(pairs, *, swiglu, limit_f32, neg_limit_f32, range_conste __all__ = [ "LOG2E", + "SituV2Consts", "batched_silu_swiglu", + "batched_situv2", "fclamp_f32", "fmin_f32", "fused_silu_swiglu_elem", + "fused_situv2_elem", "make_lds_copy_ops", "pipeline_fence", + "situv2_consts", "workgroup_barrier", ] diff --git a/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py b/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py index fea0694102..60d3ad1111 100644 --- a/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py +++ b/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py @@ -16,9 +16,12 @@ from .gemm_common_gfx1250 import ( batched_silu_swiglu, + batched_situv2, fused_silu_swiglu_elem, + fused_situv2_elem, make_lds_copy_ops, pipeline_fence, + situv2_consts, workgroup_barrier, ) from .quant_utils import ( @@ -58,6 +61,10 @@ def launch_gemm_a8w4_tdm( stage1_quant_out: Constexpr[int] = 0, quant_wmma_rep: Constexpr[int] = 1, arg_quant_scale: fx.Tensor = None, + f32_situ_beta: fx.Float32 = 1.0, + f32_situ_inv_beta: fx.Float32 = 1.0, + f32_situ_linear_beta: fx.Float32 = 1.0, + f32_situ_inv_linear_beta: fx.Float32 = 1.0, ): cache_tag = ( K, @@ -149,6 +156,10 @@ def kernel( i32_m: fx.Int32, i32_n: fx.Int32, f32_swiglu_limit: fx.Float32, + f32_situ_beta: fx.Float32, + f32_situ_inv_beta: fx.Float32, + f32_situ_linear_beta: fx.Float32, + f32_situ_inv_linear_beta: fx.Float32, ): # rocdl.disable_xdl_arb_stall() @@ -547,6 +558,19 @@ def compute_ktile(buf, prefetch_kt): STORE_N = (tile_n // 2) if stage1_act else tile_n neg_limit = fx.Float32(0.0) - f32_swiglu_limit is_swiglu = stage1_act == 2 + is_situv2 = stage1_act == 3 + # Uniform across the tile, so fold the betas once here rather than + # per element. Only materialised on the SiTUv2 path. + situ_c = ( + situv2_consts( + f32_situ_beta, + f32_situ_inv_beta, + f32_situ_linear_beta, + f32_situ_inv_linear_beta, + ) + if const_expr(is_situv2) + else None + ) oc = fx.Float16 if out_is_f16 else fx.BFloat16 # -- Activate + stage to LDS -- @@ -586,13 +610,20 @@ def compute_ktile(buf, prefetch_kt): for p in range_constexpr(4): pairs.append((acc[2 * p], acc[2 * p + 1])) - all_vals = batched_silu_swiglu( - pairs, - swiglu=is_swiglu, - limit_f32=f32_swiglu_limit, - neg_limit_f32=neg_limit, - range_constexpr=range_constexpr, - ) + if const_expr(is_situv2): + all_vals = batched_situv2( + pairs, + consts=situ_c, + range_constexpr=range_constexpr, + ) + else: + all_vals = batched_silu_swiglu( + pairs, + swiglu=is_swiglu, + limit_f32=f32_swiglu_limit, + neg_limit_f32=neg_limit, + range_constexpr=range_constexpr, + ) scale_f32, e8m0_byte = emit_amax_e8m0_native_scale( all_vals, wave_size=WAVE, dtype=MxDtype.FP8_E4M3 @@ -657,8 +688,17 @@ def compute_ktile(buf, prefetch_kt): ) ).to(fx.Float32) if const_expr(stage1_act): - hv = Vec.from_elements( - [ + if const_expr(is_situv2): + act_vals = [ + fused_situv2_elem( + acc[2 * p], + acc[2 * p + 1], + consts=situ_c, + ) + for p in range_constexpr(4) + ] + else: + act_vals = [ fused_silu_swiglu_elem( acc[2 * p], acc[2 * p + 1], @@ -667,9 +707,8 @@ def compute_ktile(buf, prefetch_kt): neg_limit_f32=neg_limit, ) for p in range_constexpr(4) - ], - fx.Float32, - ).to(oc) + ] + hv = Vec.from_elements(act_vals, fx.Float32).to(oc) lds_store_b64( stC_idx, (row_rel * STORE_N + col_rel // 2) * 2, @@ -725,6 +764,10 @@ def compute_ktile(buf, prefetch_kt): i32_m, N, f32_swiglu_limit, + f32_situ_beta, + f32_situ_inv_beta, + f32_situ_linear_beta, + f32_situ_inv_linear_beta, ).launch(grid=(m_tiles * n_tiles, 1, 1), block=(block, 1, 1), stream=stream) diff --git a/aiter/ops/flydsl/moe_common.py b/aiter/ops/flydsl/moe_common.py index 973e4fea0e..f48594d425 100644 --- a/aiter/ops/flydsl/moe_common.py +++ b/aiter/ops/flydsl/moe_common.py @@ -36,9 +36,8 @@ def apply_gate_up( ) -> torch.Tensor: """Torch reference for the stage1 gate/up activation. - ``situv2`` has no kernel on the grouped path right now -- see TODO(situv2) - in ``grouped_moe_gfx1250`` -- but the reference is kept here so the - restored kernel has something to be checked against. + ``situv2`` (Kimi-K3 ``hidden_act="situ"``) is the grouped TDM stage1 + epilogue's ``stage1_act=3``; this is what that kernel is checked against. """ lim = 7.0 if swiglu_limit is None else float(swiglu_limit) if act == "swiglu": diff --git a/op_tests/test_flydsl_grouped_gemm_gfx1250.py b/op_tests/test_flydsl_grouped_gemm_gfx1250.py index ca67d70122..a8f40b85e6 100644 --- a/op_tests/test_flydsl_grouped_gemm_gfx1250.py +++ b/op_tests/test_flydsl_grouped_gemm_gfx1250.py @@ -77,6 +77,12 @@ def parse_num_expert_activated(): SCALE_BLOCK = 32 DEFAULT_SCALE_BYTE = 127 # e8m0 byte for 2^0 = 1.0 +_ACT_BY_NAME = { + "silu": ActivationType.Silu, + "swiglu": ActivationType.Swiglu, + "situv2": ActivationType.Situv2, +} + VERIFY_TOL_A4W4 = 0.02 VERIFY_TOL_A8W4 = 0.02 # Production MoE accuracy gate (matches op_tests/test_moe_2stage.py calc_diff): @@ -670,11 +676,6 @@ def test_situv2_activation_matches_torch(): torch.testing.assert_close(actual, expected) -@pytest.mark.skip( - reason="SiTUv2 has no grouped kernel since the fused stage1 epilogue was " - "removed; the TDM path runs it as silu. See TODO(situv2) in " - "grouped_moe_gfx1250." -) def test_grouped_a4w4_situv2_matches_torch_ref(): run_moe( "a4w4", @@ -684,6 +685,18 @@ def test_grouped_a4w4_situv2_matches_torch_ref(): ) +def test_grouped_a8w4_situv2_matches_torch_ref(): + # a8w4 takes the fused stage1 quant epilogue (batched activation), which is + # a separate code path from a4w4's bf16 intermediate (element-wise). + run_moe( + "a8w4", + activation=ActivationType.Situv2, + model_dim=512, + inter_dim=512, + tol=VERIFY_TOL_A8W4, + ) + + @pytest.mark.parametrize("activation", [ActivationType.Silu, ActivationType.Swiglu]) def test_grouped_a4w4_swiglu_limit_clamps(activation): run_moe("a4w4", activation=activation, swiglu_limit=1.0) @@ -819,9 +832,7 @@ def run_csv_scenario(args) -> None: activation_override = None if args.act is not None: - activation_override = ( - ActivationType.Swiglu if args.act == "swiglu" else ActivationType.Silu - ) + activation_override = _ACT_BY_NAME[args.act] rows = [] for idx, rec in enumerate(csv_rows): @@ -981,14 +992,28 @@ def main() -> None: parser.add_argument("--iters", type=int, default=101) parser.add_argument( "--act", - choices=("silu", "swiglu"), + choices=("silu", "swiglu", "situv2"), default=None, help="stage1 activation: silu => silu(gate)*up; " - "swiglu => gpt-oss swiglu with clamp/alpha/residual. Default: swiglu " + "swiglu => gpt-oss swiglu with clamp/alpha/residual; " + "situv2 => Kimi-K3 SiTUv2 (see --situ-beta / --situ-linear-beta). " + "Default: swiglu " "for bench/verify/kernel; for --scenario csv, unset means use each " "row's act_type (pass --act to force one activation for all rows).", ) parser.add_argument("--swiglu-limit", type=float, default=7.0) + parser.add_argument( + "--situ-beta", + type=float, + default=4.0, + help="SiTUv2 gate beta (Kimi-K3 activation_situ_beta).", + ) + parser.add_argument( + "--situ-linear-beta", + type=float, + default=25.0, + help="SiTUv2 up beta (Kimi-K3 activation_situ_linear_beta).", + ) parser.add_argument( "--no-bias", action="store_true", @@ -1046,7 +1071,7 @@ def main() -> None: # sets args.tokens to a single int so run_moe reads it unchanged. token_list = args.tokens if isinstance(args.tokens, list) else [args.tokens] # None (unset) defaults to swiglu for the single-shape scenarios. - activation = ActivationType.Silu if args.act == "silu" else ActivationType.Swiglu + activation = _ACT_BY_NAME.get(args.act, ActivationType.Swiglu) rows = [] for _tok in token_list: args.tokens = _tok @@ -1066,6 +1091,8 @@ def main() -> None: tol=tol, activation=activation, swiglu_limit=args.swiglu_limit, + situ_beta=args.situ_beta, + situ_linear_beta=args.situ_linear_beta, use_bias=not args.no_bias, check_aot_cache=not args.no_check_aot_cache, raise_on_fail=False, From 72dcab488a3f0d9e1eacfefdc8b98d4a75341de4 Mon Sep 17 00:00:00 2001 From: XiaobingSuper Date: Fri, 31 Jul 2026 11:39:48 +0000 Subject: [PATCH 3/5] op_tests: cover the contiguous-M prefix scan across the block size The scan's width is set by the expert count, not the token count, so it was the one part of the grouped-MoE pipeline with no coverage on the axis that actually breaks it. Nothing else in this file varies E far enough to notice: a dropped expert does not show up as a bad number, it shows up as a row index pointing outside the contiguous buffer, and then as a fault in whichever kernel dereferences that row next. Check starts/psum/contiguous_m against a tile-aligned torch.cumsum, and separately check that every remapped route row lands inside the buffer, at E = 8/256/512/513/896/1024 -- either side of MAX_EXPERTS_PER_BLOCK, including Kimi-K3's 896. Counts come from a real unbalanced random routing rather than a uniform split, so the per-expert values differ. --- op_tests/test_flydsl_grouped_gemm_gfx1250.py | 93 ++++++++++++++++++++ 1 file changed, 93 insertions(+) diff --git a/op_tests/test_flydsl_grouped_gemm_gfx1250.py b/op_tests/test_flydsl_grouped_gemm_gfx1250.py index a8f40b85e6..dc9ef7bb7b 100644 --- a/op_tests/test_flydsl_grouped_gemm_gfx1250.py +++ b/op_tests/test_flydsl_grouped_gemm_gfx1250.py @@ -702,6 +702,99 @@ def test_grouped_a4w4_swiglu_limit_clamps(activation): run_moe("a4w4", activation=activation, swiglu_limit=1.0) +# --------------------------------------------------------------------------- +# Contiguous-M prefix scan +# +# The scan sits in front of every grouped MoE launch: it turns the per-expert +# row counts into the tile-aligned starts/psum the GEMM schedules on, and +# rewrites the route rows in place. It is also the one piece whose width is set +# by the expert count rather than the token count, so it gets its own coverage +# above and below the block size -- a wrong row here does not produce a bad +# number, it produces an out-of-bounds write in whichever kernel consumes the +# row next. +# --------------------------------------------------------------------------- +def _psum_ref(masked_m: torch.Tensor, tile_m: int): + """starts / psum / contiguous_m from a tile-aligned cumulative sum.""" + aligned = ((masked_m + tile_m - 1) // tile_m) * tile_m + inclusive = torch.cumsum(aligned.to(torch.int64), 0) + starts = inclusive - aligned + return ( + starts.to(torch.int32), + (starts + masked_m).to(torch.int32), + max(int(inclusive[-1]), tile_m), + ) + + +def _random_route_counts(experts: int, topk: int, tokens: int, seed: int = 0): + """Per-expert counts from a real (unbalanced) random routing.""" + torch.manual_seed(seed) + topk = min(topk, experts) + topk_ids = torch.stack([torch.randperm(experts)[:topk] for _ in range(tokens)]).to( + torch.int32 + ) + counts = torch.bincount(topk_ids.reshape(-1).long(), minlength=experts) + return topk_ids, counts.to(torch.int32) + + +# 512 is MAX_EXPERTS_PER_BLOCK: one thread per expert covers E up to that in a +# single pass, and everything above it needs the chunked sweep. Kimi-K3 is 896. +@pytest.mark.parametrize("experts", [8, 256, 512, 513, 896, 1024]) +def test_contiguous_psum_matches_cumsum(experts): + _require_gfx1250() + from aiter.ops.flydsl.grouped_moe_gfx1250 import contiguous_psum + + tile_m = 64 + _topk_ids, masked_m = _random_route_counts(experts, topk=16, tokens=128) + ref_starts, ref_psum, ref_total = _psum_ref(masked_m, tile_m) + + starts, psum, contiguous_m = contiguous_psum(masked_m, experts, tile_m) + torch.cuda.synchronize() + + bad = int((starts != ref_starts).sum()) + assert bad == 0, ( + f"E={experts}: {bad} experts have a wrong start, first at " + f"{int((starts != ref_starts).nonzero()[0][0])}" + ) + assert torch.equal(psum, ref_psum), f"E={experts}: psum mismatch" + assert ( + int(contiguous_m[0]) == ref_total + ), f"E={experts}: contiguous_m {int(contiguous_m[0])} != {ref_total}" + + +@pytest.mark.parametrize("experts", [8, 256, 512, 513, 896, 1024]) +def test_contiguous_psum_remap_rows_stay_in_bounds(experts): + """The remap is what the MoE actually calls; an unscanned expert lands here + as a row index pointing outside the contiguous buffer.""" + _require_gfx1250() + from aiter.ops.flydsl.grouped_moe_gfx1250 import contiguous_psum_remap + + tile_m, topk, tokens = 64, 16, 128 + topk_ids, masked_m = _random_route_counts(experts, topk, tokens) + ref_starts, _ref_psum, ref_total = _psum_ref(masked_m, tile_m) + + # Masked layout: row = expert * max_m + slot, which is what + # flydsl_moe_topids_to_rows produces and the remap folds down. + flat = topk_ids.reshape(-1) + max_m = max(tile_m, ((flat.numel() + tile_m - 1) // tile_m) * tile_m) + slot = torch.zeros(experts, dtype=torch.int64) + rows = torch.empty(flat.numel(), dtype=torch.int32) + for i, e in enumerate(flat.tolist()): + rows[i] = e * max_m + int(slot[e]) + slot[e] += 1 + + remapped = rows.clone() + contiguous_psum_remap(masked_m, remapped, experts, max_m, tile_m) + torch.cuda.synchronize() + + expected = ref_starts[flat.long()].long() + (rows.long() - flat.long() * max_m) + assert torch.equal(remapped.long(), expected), f"E={experts}: row remap mismatch" + oob = int((remapped >= ref_total).sum()) + assert oob == 0, ( + f"E={experts}: {oob} remapped rows land outside the contiguous buffer " + f"(bound {ref_total}, max row {int(remapped.max())})" + ) + + # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- From 4f0df77044b1d4675e8f39e78312a3678b7d5a51 Mon Sep 17 00:00:00 2001 From: XiaobingSuper Date: Fri, 31 Jul 2026 12:08:16 +0000 Subject: [PATCH 4/5] gfx1250: only validate the SiTUv2 betas when SiTUv2 is what runs The bounds check sat at the top of flydsl_grouped_gemm_a8w4_masked, so it applied to silu and swiglu launches too -- where the betas are ignored and default to 1.0. A caller that passed a beta of 0 alongside a non-SiTUv2 activation would have been rejected for a parameter the kernel never reads. Gate it on stage1_act == 3. --- aiter/ops/flydsl/batched_gemm_mxfp4.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/aiter/ops/flydsl/batched_gemm_mxfp4.py b/aiter/ops/flydsl/batched_gemm_mxfp4.py index 0acf83f7fd..c8db04ddff 100644 --- a/aiter/ops/flydsl/batched_gemm_mxfp4.py +++ b/aiter/ops/flydsl/batched_gemm_mxfp4.py @@ -83,10 +83,13 @@ def flydsl_grouped_gemm_a8w4_masked( if stream is None: stream = torch.cuda.current_stream() - if float(situ_beta) <= 0.0: - raise ValueError(f"situ_beta must be > 0, got {situ_beta!r}") - if float(situ_linear_beta) <= 0.0: - raise ValueError(f"situ_linear_beta must be > 0, got {situ_linear_beta!r}") + # Only meaningful for SiTUv2; the betas are ignored by every other epilogue, + # so do not let them reject a silu/swiglu launch. + if stage1_act == 3: + if float(situ_beta) <= 0.0: + raise ValueError(f"situ_beta must be > 0, got {situ_beta!r}") + if float(situ_linear_beta) <= 0.0: + raise ValueError(f"situ_linear_beta must be > 0, got {situ_linear_beta!r}") nb = min(num_buffers, max(1, K // tile_k)) has_bias = 1 if bias is not None else 0 bias_ptr = ptr_arg(bias) if bias is not None else ptr_arg(a) From 01f8c9fb442e3c34dc517ff33876de763c45d20d Mon Sep 17 00:00:00 2001 From: XiaobingSuper Date: Mon, 3 Aug 2026 04:39:39 +0000 Subject: [PATCH 5/5] gfx1250: take the SiTUv2 reciprocals in-kernel, and guard the one scan still capped at 512 Review feedback on #4482. Drop f32_situ_inv_beta / f32_situ_inv_linear_beta from the TDM kernel and let situv2_consts() take both reciprocals with v_rcp_f32. Both are uniform across the tile, so this is two extra VALU ops per kernel, hoisted out of the inner loop, against two fewer kernel args and no way for a caller to pass a beta and a reciprocal that disagree. Verified numerically identical on the Kimi-K3 betas: a4w4 logits_diff 9.0759e-06 / rel_l2 4.2605e-03 and a8w4 4.4326e-06 / 2.9774e-03 both match the host-reciprocal build to every printed digit (beta=4.0 is a power of two so its rcp is exact; linear_beta=25.0 is not, and its ~1 ulp lands far below the MXFP4 quantisation it feeds). Rewrite the chunked-scan comment to say what it is actually guarding: one Hillis-Steele pass covers one expert per lane, so the old single-pass scan left starts/psum unwritten for every expert past 512, which is how K3's 896 reached the GEMM as garbage offsets and faulted. Sweeping the other single-block scans for the same cap: the route-quant-scatter prefix sum is single-thread serial over E (no cap) and moe_g2l_lut is gated at _G2L_MAX_N with a torch fallback, both fine. moe_route_psum_fused is genuinely capped -- its LDS route counter is one slot per expert, so E>512 needs a wider allocation, not a carry -- and _FUSED_ROUTE_PSUM_MAX_EXPERTS was defined but never enforced. Raise instead of silently dropping experts. The NUMEL companion is left advisory: that sweep is grid-stride, so a larger count is correct, just not worth fusing. Co-Authored-By: Claude Opus 5 --- aiter/ops/flydsl/batched_gemm_mxfp4.py | 5 +-- aiter/ops/flydsl/grouped_moe_gfx1250.py | 16 +++++++- .../ops/flydsl/kernels/gemm_common_gfx1250.py | 16 +++++--- .../ops/flydsl/kernels/moe_contiguous_psum.py | 38 ++++++++++++++----- .../kernels/mxfp4_preshuffle_gfx1250_tdm.py | 13 +------ 5 files changed, 54 insertions(+), 34 deletions(-) diff --git a/aiter/ops/flydsl/batched_gemm_mxfp4.py b/aiter/ops/flydsl/batched_gemm_mxfp4.py index c8db04ddff..233d6a711b 100644 --- a/aiter/ops/flydsl/batched_gemm_mxfp4.py +++ b/aiter/ops/flydsl/batched_gemm_mxfp4.py @@ -70,8 +70,7 @@ def flydsl_grouped_gemm_a8w4_masked( ``stage1_act`` selects the stage1 epilogue: 0 none, 1 silu, 2 swiglu, 3 SiTUv2 (``situ_beta`` / ``situ_linear_beta``, the Kimi-K3 activation). The betas are runtime kernel arguments, so all SiTUv2 shapes share one - compiled kernel; their reciprocals are taken here rather than by an - in-kernel v_rcp_f32 so the fold stays exact. + compiled kernel. When ``stage1_quant_out=1`` (fp8), the epilogue fuses the activation + MX fp8 quantization + e8m0 scale preshuffle into the kernel. ``out`` receives @@ -126,9 +125,7 @@ def flydsl_grouped_gemm_a8w4_masked( quant_wmma_rep, quant_scale_tensor, float(situ_beta), - 1.0 / float(situ_beta), float(situ_linear_beta), - 1.0 / float(situ_linear_beta), ) return out diff --git a/aiter/ops/flydsl/grouped_moe_gfx1250.py b/aiter/ops/flydsl/grouped_moe_gfx1250.py index 722f151bb2..eda290d693 100644 --- a/aiter/ops/flydsl/grouped_moe_gfx1250.py +++ b/aiter/ops/flydsl/grouped_moe_gfx1250.py @@ -1094,8 +1094,10 @@ def _get_compiled_route_psum_fused(): return build_moe_route_psum_fused_module() -# One workgroup handles every route, so the fused kernel only applies while the -# route count fits a single block's grid-stride sweep and E fits the scan. +# One workgroup handles every route. NUMEL is advisory -- the route sweep is +# grid-stride, so a larger count is correct but stops being worth fusing. +# EXPERTS is a hard limit, enforced below: the scan and the LDS route counter +# are both one slot per lane. _FUSED_ROUTE_PSUM_MAX_NUMEL = 4096 _FUSED_ROUTE_PSUM_MAX_EXPERTS = 512 @@ -1116,6 +1118,16 @@ def fused_route_psum_remap( token_num, topk = topk_ids.shape numel = token_num * topk experts = int(experts) + # Unlike contiguous_psum/_remap, this kernel's scan is still single-pass: + # its LDS route counter is one slot per expert, so widening E needs a bigger + # allocation, not just a carry. Fail loudly rather than silently drop the + # experts past the block, which is the bug the chunked scan fixed there. + if experts > _FUSED_ROUTE_PSUM_MAX_EXPERTS: + raise ValueError( + f"fused_route_psum_remap supports at most " + f"{_FUSED_ROUTE_PSUM_MAX_EXPERTS} experts, got {experts}; " + f"use flydsl_moe_topids_to_rows + contiguous_psum_remap instead" + ) topids_to_rows = torch.empty(numel, dtype=torch.int32, device=device) masked_m = torch.empty(experts, dtype=torch.int32, device=device) starts = torch.empty(experts, dtype=torch.int32, device=device) diff --git a/aiter/ops/flydsl/kernels/gemm_common_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_common_gfx1250.py index 2343380dff..bbd5212251 100644 --- a/aiter/ops/flydsl/kernels/gemm_common_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_common_gfx1250.py @@ -122,21 +122,25 @@ def _tanh_f32(x, tanh_mul): SituV2Consts = namedtuple("SituV2Consts", "beta gate_tanh_mul linear_beta up_tanh_mul") -def situv2_consts(beta, inv_beta, linear_beta, inv_linear_beta): +def situv2_consts(beta, linear_beta): """Fold the SiTUv2 betas into the per-element multipliers, once per kernel. - ``inv_beta`` / ``inv_linear_beta`` are the host-computed reciprocals (exact, - unlike an in-kernel v_rcp_f32). Everything here is uniform across the tile, - so hoisting it keeps the inner loop at 3 exp2 + 3 rcp per element. + The two reciprocals are taken here with v_rcp_f32 rather than passed in from + the host: both are uniform across the tile, so this is two extra VALU ops per + kernel, hoisted out of the inner loop, in exchange for two fewer kernel args + and no way for a caller to hand in a beta and a reciprocal that disagree. + v_rcp_f32's ~1 ulp sits far below the MXFP4 quantisation this feeds. + + Hoisting keeps the inner loop at 3 exp2 + 3 rcp per element. """ import flydsl.expr as _fx neg_two_log2e = _fx.Float32(-2.0 * LOG2E) return SituV2Consts( beta=beta, - gate_tanh_mul=neg_two_log2e * inv_beta, + gate_tanh_mul=neg_two_log2e * _fx.Float32(rocdl.rcp(T.f32, _raw(beta))), linear_beta=linear_beta, - up_tanh_mul=neg_two_log2e * inv_linear_beta, + up_tanh_mul=neg_two_log2e * _fx.Float32(rocdl.rcp(T.f32, _raw(linear_beta))), ) diff --git a/aiter/ops/flydsl/kernels/moe_contiguous_psum.py b/aiter/ops/flydsl/kernels/moe_contiguous_psum.py index e6b0493c9a..c8b1e586cd 100644 --- a/aiter/ops/flydsl/kernels/moe_contiguous_psum.py +++ b/aiter/ops/flydsl/kernels/moe_contiguous_psum.py @@ -109,11 +109,20 @@ def psum_kernel( _lds_store(carry, fx.Int32(0), 0) gpu.barrier() - # E is not bounded by the block width: sweep the experts in block-sized - # chunks, offsetting each chunk by the running total of the ones before - # it. That total lives in LDS (``carry``) because a register would not - # survive the runtime chunk loop. Lanes past ``experts`` feed 0 into the - # scan so the last lane still holds the chunk total, and write nothing. + # One Hillis-Steele scan spans exactly one thread per lane, so a single + # pass covers at most MAX_EXPERTS_PER_BLOCK experts -- it used to be the + # whole kernel, which silently left starts/psum unwritten for every + # expert past 512 (Kimi-K3 has 896: garbage offsets, then a memory fault + # in the GEMM that indexes with them). + # + # So sweep E in block-sized chunks instead. Each chunk scans as before + # and then adds ``carry``, the tile-aligned total of all chunks already + # scanned, which is what makes the per-chunk scans one continuous prefix + # sum. ``carry`` has to be LDS, not a register: it is produced by lane 0 + # and consumed by all of them on the next iteration. + # + # Lanes past ``experts`` feed 0 into the scan -- they keep the last lane + # holding the true chunk total, and write no output. for base in range(0, experts, MAX_EXPERTS_PER_BLOCK): e = fx.Uint32(base) + tid in_expert = e < fx.Uint32(experts) @@ -230,11 +239,20 @@ def psum_remap_kernel( _lds_store(carry, fx.Int32(0), 0) gpu.barrier() - # E is not bounded by the block width: sweep the experts in block-sized - # chunks, offsetting each chunk by the running total of the ones before - # it. That total lives in LDS (``carry``) because a register would not - # survive the runtime chunk loop. Lanes past ``experts`` feed 0 into the - # scan so the last lane still holds the chunk total, and write nothing. + # One Hillis-Steele scan spans exactly one thread per lane, so a single + # pass covers at most MAX_EXPERTS_PER_BLOCK experts -- it used to be the + # whole kernel, which silently left starts/psum unwritten for every + # expert past 512 (Kimi-K3 has 896: garbage offsets, then a memory fault + # in the GEMM that indexes with them). + # + # So sweep E in block-sized chunks instead. Each chunk scans as before + # and then adds ``carry``, the tile-aligned total of all chunks already + # scanned, which is what makes the per-chunk scans one continuous prefix + # sum. ``carry`` has to be LDS, not a register: it is produced by lane 0 + # and consumed by all of them on the next iteration. + # + # Lanes past ``experts`` feed 0 into the scan -- they keep the last lane + # holding the true chunk total, and write no output. for base in range(0, experts, MAX_EXPERTS_PER_BLOCK): e = fx.Uint32(base) + tid in_expert = e < fx.Uint32(experts) diff --git a/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py b/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py index 60d3ad1111..9cdbd5184e 100644 --- a/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py +++ b/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py @@ -62,9 +62,7 @@ def launch_gemm_a8w4_tdm( quant_wmma_rep: Constexpr[int] = 1, arg_quant_scale: fx.Tensor = None, f32_situ_beta: fx.Float32 = 1.0, - f32_situ_inv_beta: fx.Float32 = 1.0, f32_situ_linear_beta: fx.Float32 = 1.0, - f32_situ_inv_linear_beta: fx.Float32 = 1.0, ): cache_tag = ( K, @@ -157,9 +155,7 @@ def kernel( i32_n: fx.Int32, f32_swiglu_limit: fx.Float32, f32_situ_beta: fx.Float32, - f32_situ_inv_beta: fx.Float32, f32_situ_linear_beta: fx.Float32, - f32_situ_inv_linear_beta: fx.Float32, ): # rocdl.disable_xdl_arb_stall() @@ -562,12 +558,7 @@ def compute_ktile(buf, prefetch_kt): # Uniform across the tile, so fold the betas once here rather than # per element. Only materialised on the SiTUv2 path. situ_c = ( - situv2_consts( - f32_situ_beta, - f32_situ_inv_beta, - f32_situ_linear_beta, - f32_situ_inv_linear_beta, - ) + situv2_consts(f32_situ_beta, f32_situ_linear_beta) if const_expr(is_situv2) else None ) @@ -765,9 +756,7 @@ def compute_ktile(buf, prefetch_kt): N, f32_swiglu_limit, f32_situ_beta, - f32_situ_inv_beta, f32_situ_linear_beta, - f32_situ_inv_linear_beta, ).launch(grid=(m_tiles * n_tiles, 1, 1), block=(block, 1, 1), stream=stream)