[AMD/MI300A] mxfp4 fused-MoE: support expert parallelism (ep_size>1) on gfx942 (DeepSeek-V4)

[xzyaoi]

Context

Bringing up DeepSeek-V4 (model_type: deepseek_v4 — FP4/mxfp4 routed experts + FP8 for the rest) on AMD MI300A (gfx942, ROCm 7.2, torch 2.11). V4 has 256 experts (Flash) / 384 experts (Pro), top-6 — so the experts must be sharded across GPUs/nodes (expert parallelism) to fit (V4-Pro is 806 GB; even Flash's experts don't fit one 128 GB GPU replicated).

Finding (single-node probe, MI300A, 2026-06-11)

The gfx942 mxfp4 MoE kernel works — but only without expert parallelism:

• ep_size = 4 (--enable-expert-parallel) → construction fails at backend selection:

  RuntimeError: No MoE backend available for gfx94/mxfp4.
    Tried: gluon_kernel:unsupported, triton_kernel:unsupported
    (tokenspeed/runtime/layers/moe/core/selector.py:192)
  

• ep_size = 1 (TP only) → Mxfp4TritonKernelBackend selects and allocates mxfp4 weights (mxfp4/triton_kernel.py:153 create_layer_weights → mxfp4/weights.py:66 create_mxfp4_weights); it only then OOMs because 806 GB can't fit one 128 GB GPU. So the kernel itself is fine on gfx942.

Root cause

tokenspeed/runtime/layers/moe/backends/mxfp4/{gluon,triton}_kernel.py supports() ends with:

return spec.ep_size <= 1 and spec.activation in {"silu", "swiglu"}

i.e. the AMD mxfp4 fused-MoE (w-mxfp4 / a-fp8) is gated to ep_size ≤ 1. There is no mxfp4 MoE path under expert parallelism, so V4's experts can't be sharded → the model can't be served on MI300A.

Ask

Extend the AMD mxfp4 fused-MoE (gluon/triton) to support ep_size > 1: per-rank expert subset + masked-compute + all-reduce, mirroring the bf16 / INT4-W4A16 EP path that already serves Kimi-K2.6 on 2×MI300A. The ep_size <= 1 gate is tokenspeed-side; the kernel may already accept a per-rank expert subset (please verify) — the work is lifting the gate + wiring the EP masked-compute (plus any kernel changes if the layout assumes all experts are local).

Repro

jobs/serve-v4pro-1node-probe.sbatch (ep=4 → the selector error) and serve-v4pro-1node-probe2.sbatch (ep=1 → selects + OOM) on beverin. V4-Pro weights staged at infra01/hf_models/.../DeepSeek-V4-Pro.

HW: MI300A gfx942 / ROCm 7.2 / torch 2.11.0+rocm7.2. Part of the DeepSeek-V4-on-MI300A bring-up (tracking issue to follow).

[xzyaoi]

Shipped the host-side expert-parallelism wiring for the portable INT4-W4A16 mirror of the AMD quantized fused-MoE in #30 (draft).

What it adds

• moe_align_block_size_ep(topk_ids, block_size, num_experts, expert_map): remaps global routing ids to per-rank local weight rows via expert_map (-1 = not on this rank), drops non-local slots, and reuses the existing per-expert block builder over the E_local local experts. expert_ids then indexes the rank-local [E_local, N, K//8] weights directly.
• fused_moe_int4_w4a16(..., expert_map=None): opt-in EP kwarg on the public op + reference + Triton backends. When given, packed/scale are the rank-local slice and the op returns this rank's partial [M, N] output; summing the partials across the EP group (the production all-reduce) reconstructs the full dense result. expert_map=None is byte-for-byte the prior behavior.

Confirmed the issue's hypothesis: the kernel already accepts a per-rank expert subset — the grouped GEMM tiles by expert, filters via expert_ids, and gathers tokens by sorted_token_ids // top_k. So no kernel-body change was needed; the ep_size <= 1 gate was purely host-side, and the missing piece was the dispatch builder + plumbing.

Validation (acceptance: sum_over_ranks(partial[r]) == full non-EP output, bf16 2e-2):

• Local CPU (TRITON_INTERPRET=1, fp32 3e-3): full suite 105 passed, 3 skipped; new EP test 11/11 (slot-coverage, sum(partials)==dense for mul_routed ∈ {False,True}, ep ∈ {2,3,4}, identity-map==non-EP, empty-rank→zero).
• beverin MI300A (gfx942, bf16, job 381967): EP test 11/11 PASS at decode (E=48, top_k=8, ep=4) + prefill shapes, all within 2e-2; INT4 MoE suite 12/12 (no regression).

Caveats

• This is the single-GPU kernel + host dispatch. The cross-rank all-reduce of the partials is a separate collective the serving layer owns (xkernels.ops.comm, Topology-aware hierarchical all-reduce for DP-attention MoE serving (2-node MI300A) #12). Lifting the tokenspeed ep_size <= 1 gate to call this path is the framework-side follow-up.
• No perf claim — EP changes which experts a rank computes, not the per-expert GEMM. The win is fitting V4's experts across ranks, which the gate previously blocked.

Details: docs/issue-26-mxfp4-moe-ep.md. PR: #30

[xzyaoi]

Closed by #30 — adds the host-side expert-parallelism wiring (moe_align_block_size_ep + per-rank expert-subset / masked-compute dispatch) that lifts the ep_size<=1 gate for the quantized fused-MoE GEMM. The grouped GEMM already tiles by expert and filters per-block via an expert_ids sentinel, so no kernel-body change was needed. Merged 2026-06-11.