gfx1250: fix the grouped-MoE expert scan above 512 experts, and compute SiTUv2 instead of SiLU

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
 

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,17 +1014,15 @@ 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,
         )
 
     # Only the felix TDM grouped path is kept; the previous non-TDM grouped
     # 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)

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",
 ]