diff --git a/aiter/ops/flydsl/batched_gemm_mxfp4.py b/aiter/ops/flydsl/batched_gemm_mxfp4.py index 3f546349bb..233d6a711b 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,28 @@ 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. + + 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() + # 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) @@ -110,6 +124,8 @@ def flydsl_grouped_gemm_a8w4_masked( stage1_quant_out, quant_wmma_rep, quant_scale_tensor, + float(situ_beta), + 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..eda290d693 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. @@ -1080,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 @@ -1102,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 24f0bb01b5..bbd5212251 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,125 @@ 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, linear_beta): + """Fold the SiTUv2 betas into the per-element multipliers, once per kernel. + + 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 * _fx.Float32(rocdl.rcp(T.f32, _raw(beta))), + linear_beta=linear_beta, + up_tanh_mul=neg_two_log2e * _fx.Float32(rocdl.rcp(T.f32, _raw(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 +269,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/moe_contiguous_psum.py b/aiter/ops/flydsl/kernels/moe_contiguous_psum.py index ca84b12ffa..c8b1e586cd 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,78 @@ 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: + # 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) + 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 +226,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 +234,71 @@ 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 + # 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) + 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() diff --git a/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py b/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py index fea0694102..9cdbd5184e 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,8 @@ 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_linear_beta: fx.Float32 = 1.0, ): cache_tag = ( K, @@ -149,6 +154,8 @@ def kernel( i32_m: fx.Int32, i32_n: fx.Int32, f32_swiglu_limit: fx.Float32, + f32_situ_beta: fx.Float32, + f32_situ_linear_beta: fx.Float32, ): # rocdl.disable_xdl_arb_stall() @@ -547,6 +554,14 @@ 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_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 +601,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 +679,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 +698,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 +755,8 @@ def compute_ktile(buf, prefetch_kt): i32_m, N, f32_swiglu_limit, + f32_situ_beta, + f32_situ_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..dc9ef7bb7b 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,11 +685,116 @@ 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) +# --------------------------------------------------------------------------- +# 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 # --------------------------------------------------------------------------- @@ -819,9 +925,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 +1085,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 +1164,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 +1184,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,