[Triton/Gluon] [fix] moe_gemm_mxfp8: zero-init output, graph-safe routing, skip redundant copy

aiter/ops/triton/README.md

@@ -358,6 +358,7 @@ deleted. **Do not re-implement a shuffle** (the telltale is a local
 | ------------------------ | ----------------------------------------------------------------------- |
 | `shuffle_weight(x, ...)` | Arch-aware weight preshuffle: gfx1250 WMMA/TDM path, otherwise delegates to `aiter.ops.shuffle.shuffle_weight` |
 | `moe_weight_decode_view(w)` | Zero-copy `(E, N, K)` → decode view sharing storage (renamed from `moe_weight_gfx1250_decode_view`) |
+| `moe_weight_kn(w)`       | `(E, N, K)` → K-major storage for `_moe_gemm_a8w8`; free view when the caller already stores K-major |
 | `shuffle_scale_gemm` / `unshuffle_scale_gemm` | GEMM MX-scale tiles — gfx950 `(32, 8)`, gfx1250 `(16, 4)` |
 | `shuffle_scale_moe`      | MoE MX scales (a8w4/a8w8/a16w4/a4w4); `return_layout=True` also returns the `SWIZZLE_MX_SCALE` label (`CDNA4_SCALE`/`GFX1250_SCALE`); no-op on arches without a native layout (e.g. gfx942) |
 | `shuffle_scale_batched`  | FP4 blockscale16 batched scales, arch-independent                        |

aiter/ops/triton/moe/moe_gemm_mxfp8.py

@@ -22,6 +22,7 @@
 from aiter.ops.triton._triton_kernels.moe.moe_op_gemm_a8w8 import _moe_gemm_a8w8
 from aiter.ops.triton.moe.moe_utils import group_sizes_to_expt_tensors
 from aiter.ops.triton.utils.logger import AiterTritonLogger
+from aiter.ops.triton.utils.shuffle import moe_weight_kn
 from aiter.ops.triton.utils.tuned_config_utils import get_tuned_kernel_config
 
 __all__ = ["moe_gemm_mxfp8"]
@@ -78,7 +79,10 @@ def moe_gemm_mxfp8(
         K % quant_block_size == 0
     ), f"K ({K}) must be divisible by quant_block_size ({quant_block_size})"
 
-    out = torch.empty(total_tokens, N, dtype=out_dtype, device=lhs.device)
+    # Rows excluded from group_sizes are never written; zero-fill so they do
+    # not surface stale allocator contents. A NaN there survives a caller's
+    # `* 0` mask (NaN * 0 == NaN) and spreads through the top-k reduction.
+    out = torch.zeros(total_tokens, N, dtype=out_dtype, device=lhs.device)
     if total_tokens == 0:
         return out
 
@@ -98,16 +102,18 @@ def moe_gemm_mxfp8(
     BLOCK_N = cfg.kwargs["BLOCK_N"]
 
     expt_hist, expt_offs, expt_offs_sum, expt_data, grid_m = (
-        group_sizes_to_expt_tensors(group_sizes, BLOCK_M)
+        group_sizes_to_expt_tensors(group_sizes, BLOCK_M, total_tokens)
     )
     if grid_m == 0:
         return out
 
     grid_n = triton.cdiv(N, BLOCK_N)
 
-    # Permute rhs/w_scale from NK to KN layout as required by _moe_gemm_a8w8.
-    rhs_kn = rhs.permute(0, 2, 1).contiguous()
-    w_scale_kn = w_scale.permute(0, 2, 1).contiguous()
+    # _moe_gemm_a8w8 wants KN. Expert weights are static after load, so a
+    # caller storing them K-major already gets a view instead of a per-call
+    # copy of the whole tensor (2.4-7.9 ms on gfx942 at Hy4 shapes).
+    rhs_kn = moe_weight_kn(rhs)
+    w_scale_kn = moe_weight_kn(w_scale)
     bias_stride = N if bias is not None else 0
 
     _moe_gemm_a8w8[(grid_m * grid_n,)](

aiter/ops/triton/moe/moe_utils.py

@@ -87,10 +87,16 @@ def build_block_mapping(
 
 
 def group_sizes_to_expt_tensors(
-    group_sizes: torch.Tensor, block_m: int
+    group_sizes: torch.Tensor, block_m: int, total_tokens: int | None = None
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]:
     """Convert group_sizes to ExptHist/ExptOffs/ExptData for _moe_gemm_a8w8.
+
     Tokens must be in contiguous expert order (no gather needed).
+
+    Passing ``total_tokens`` makes this CUDA-graph capturable: the block count
+    is then bounded statically instead of read back with ``.item()``, and the
+    padding slots carry ``ExptData = -1``, which the kernel skips. Omitting it
+    keeps the exact block count and the host sync.
     """
     E = group_sizes.shape[0]
     device = group_sizes.device
@@ -104,7 +110,12 @@ def group_sizes_to_expt_tensors(
 
     # Build ExptData = (block_id << 16) | expert_id using tensor ops (no per-expert CPU sync).
     n_blocks_per_expert = (gs + block_m - 1) // block_m  # cdiv without .item()
-    total_blocks = int(n_blocks_per_expert.sum().item())  # single sync at end
+    if total_tokens is None:
+        total_blocks = int(n_blocks_per_expert.sum().item())  # syncs; see docstring
+    else:
+        # Each expert wastes at most one partial block, so this bounds the real
+        # count without a sync. Surplus slots are stamped -1 below.
+        total_blocks = (total_tokens + block_m - 1) // block_m + E
 
     if total_blocks == 0:
         return (
@@ -124,5 +135,7 @@ def group_sizes_to_expt_tensors(
     )
     local_block_id = block_ids - block_cumsum[expert_for_block]
 
-    expt_data = ((local_block_id << 16) | expert_for_block).to(torch.int32)
+    packed = ((local_block_id << 16) | expert_for_block).to(torch.int32)
+    valid = local_block_id < n_blocks_per_expert[expert_for_block]
+    expt_data = torch.where(valid, packed, torch.full_like(packed, -1))
     return expt_hist, expt_offs, expt_offs_sum, expt_data, total_blocks

aiter/ops/triton/utils/shuffle.py

@@ -79,6 +79,18 @@ def moe_weight_decode_view(w: torch.Tensor) -> torch.Tensor:
     return w_u8.view(E, N // 16, K * 16).transpose(-1, -2)
 
 
+def moe_weight_kn(w: torch.Tensor) -> torch.Tensor:
+    """``(E, N, K)`` -> K-major storage, same logical shape.
+
+    ``_moe_gemm_a8w8`` reads the weight and its MX scales K-major. Callers that
+    already store them that way (converting once at load time rather than on
+    every call) get a free view; anything else is copied as before.
+    """
+    w_kn = w.permute(0, 2, 1)
+    already_kn = w.stride(1) == 1 and w.stride(2) == w.shape[1]
+    return w_kn if already_kn else w_kn.contiguous()
+
+
 # =============================================================================
 # SCALES
 # =============================================================================

op_tests/op_benchmarks/triton/bench_moe_gemm_mxfp8.py

@@ -15,12 +15,16 @@
 from aiter.ops.triton.moe.moe_gemm_mxfp8 import moe_gemm_mxfp8
 from op_tests.op_benchmarks.triton.utils.benchmark_utils import get_caller_name_no_ext
 
-# DSv4-style MoE shapes: (E, N, K, tokens_per_expert)
+# MoE shapes: (E, N, K, tokens_per_expert)
 _SHAPES = [
     (8, 7168, 2048, 64, "DSv4-decode"),
     (8, 7168, 2048, 256, "DSv4-prefill-small"),
     (8, 7168, 2048, 1024, "DSv4-prefill-large"),
     (64, 2048, 7168, 64, "DSv4-E64-decode"),
+    # Hy4-preview-FP8 gemm1 as sharded at TP=8: 256 experts / 8 ranks,
+    # hidden 6144, moe_intermediate 2048.
+    (32, 2048, 6144, 1, "Hy4-TP8-decode"),
+    (32, 2048, 6144, 16, "Hy4-TP8-prefill"),
 ]
 _QBS = 32
 

op_tests/triton_tests/moe/test_moe_gemm_mxfp8.py

@@ -14,15 +14,22 @@
 # E8M0 biased byte → FP32 scale: 2^(b - 127)
 _E8M0_BIAS = 127
 
+# Zero-token experts and partial blocks.
+_RAGGED_GROUPS = (0, 7, 64, 1, 100, 0, 84, 0)
+
 
 def _e8m0_to_fp32(b: torch.Tensor) -> torch.Tensor:
     return (b.to(torch.int32) << 23).view(torch.float32)
 
 
-def _make_mxfp8_inputs(E, N, K, qbs, dtype=torch.float8_e4m3fnuz, device="cuda"):
+def _make_mxfp8_inputs(
+    E, N, K, qbs, group_sizes=None, dtype=torch.float8_e4m3fnuz, device="cuda"
+):
     torch.manual_seed(0)
-    total_tokens = E * 64  # 64 tokens per expert
-    group_sizes = torch.full((E,), 64, dtype=torch.int32, device=device)
+    if group_sizes is None:
+        group_sizes = (64,) * E  # 64 tokens per expert
+    total_tokens = sum(group_sizes)
+    group_sizes = torch.tensor(group_sizes, dtype=torch.int32, device=device)
 
     lhs = torch.randint(-3, 4, (total_tokens, K), dtype=torch.int8, device=device).to(
         dtype
@@ -63,18 +70,33 @@ def _mxfp8_reference(lhs, rhs, x_scale, w_scale, group_sizes, qbs, out_dtype):
     return out
 
 
-@pytest.mark.parametrize("E, N, K", [(4, 256, 128), (8, 128, 256)])
+def _assert_matches_reference(out, ref):
+    scale = ref.abs().max().clamp_min(1e-4)
+    err = (out.float() - ref.float()).abs().max()
+    assert err <= 0.2 * scale, f"max err {err:.4f} > 0.2 * {scale:.4f}"
+
+
+@pytest.mark.parametrize(
+    "E, N, K, groups",
+    [(4, 256, 128, None), (8, 128, 256, None), (8, 128, 256, _RAGGED_GROUPS)],
+)
+@pytest.mark.parametrize("kn_weights", [False, True])
 @pytest.mark.parametrize("qbs", [32])
-def test_moe_gemm_mxfp8_correctness(E, N, K, qbs):
-    lhs, rhs, x_scale, w_scale, group_sizes = _make_mxfp8_inputs(E, N, K, qbs)
+def test_moe_gemm_mxfp8_correctness(E, N, K, groups, kn_weights, qbs):
+    """Match the reference for uniform and ragged groups, NK and KN weights."""
+    lhs, rhs, x_scale, w_scale, group_sizes = _make_mxfp8_inputs(
+        E, N, K, qbs, group_sizes=groups
+    )
+    if kn_weights:
+        # Same logical [E, N, K] shape, K-major storage: the wrapper takes a view.
+        rhs = rhs.transpose(1, 2).contiguous().transpose(1, 2)
+        w_scale = w_scale.transpose(1, 2).contiguous().transpose(1, 2)
     out = moe_gemm_mxfp8(lhs, rhs, x_scale, w_scale, group_sizes, quant_block_size=qbs)
     ref = _mxfp8_reference(lhs, rhs, x_scale, w_scale, group_sizes, qbs, torch.bfloat16)
 
     assert out.shape == ref.shape
     assert out.dtype == torch.bfloat16
-    scale = ref.abs().max().clamp_min(1e-4)
-    err = (out.float() - ref.float()).abs().max()
-    assert err <= 0.2 * scale, f"max err {err:.4f} > 0.2 * {scale:.4f}"
+    _assert_matches_reference(out, ref)
 
 
 @pytest.mark.parametrize("E, N, K", [(4, 256, 128), (8, 128, 256)])
@@ -107,3 +129,40 @@ def test_moe_gemm_mxfp8_empty_tokens():
     group_sizes = torch.zeros(E, dtype=torch.int32, device="cuda")
     out = moe_gemm_mxfp8(lhs, rhs, x_scale, w_scale, group_sizes, quant_block_size=qbs)
     assert out.shape == (0, N)
+
+
+def test_moe_gemm_mxfp8_unwritten_rows_are_zero():
+    """Rows excluded from group_sizes read as 0, not stale allocator memory."""
+    E, N, K, qbs = 4, 256, 128, 32
+    lhs, rhs, x_scale, w_scale, _ = _make_mxfp8_inputs(E, N, K, qbs)
+    covered = lhs.shape[0] // 2
+    group_sizes = torch.zeros(E, dtype=torch.int32, device="cuda")
+    group_sizes[0] = covered
+
+    # Free a NaN block of exactly the output's size so the caching allocator
+    # hands it back for `out`; an uninitialized buffer then shows up as NaN.
+    poison = torch.full(
+        (lhs.shape[0], N), float("nan"), dtype=torch.bfloat16, device="cuda"
+    )
+    del poison
+
+    out = moe_gemm_mxfp8(lhs, rhs, x_scale, w_scale, group_sizes, quant_block_size=qbs)
+    tail = out[covered:]
+    torch.testing.assert_close(tail, torch.zeros_like(tail), rtol=0, atol=0)
+
+
+def test_moe_gemm_mxfp8_cuda_graph_capture():
+    """The call is capturable: no device-to-host sync on the host path."""
+    E, N, K, qbs = 4, 256, 128, 32
+    args = _make_mxfp8_inputs(E, N, K, qbs)
+    moe_gemm_mxfp8(*args, quant_block_size=qbs)  # JIT outside the capture
+    torch.cuda.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = moe_gemm_mxfp8(*args, quant_block_size=qbs)
+    graph.replay()
+    torch.cuda.synchronize()
+
+    ref = _mxfp8_reference(*args, qbs, torch.bfloat16)
+    _assert_matches_reference(captured, ref)