fix(moe_ep): fix in_kernel_fc2_reduce livelock on zero-token launches (MXFP8 + NVFP4)

flashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/mxfp8.py

@@ -867,11 +867,34 @@ def get_symm_buffer_for_mxfp8_mega_moe(
             topk=num_topk,
             max_tokens=num_max_tokens,
         )
+    # `in_kernel_fc2_reduce` is a caller-owned correctness choice; see the
+    # NVFP4 factory. The MXFP8 kernel rejects ikr together with dispatch-warp
+    # token-back, so a tuned `token_back_mode` that isn't "epi_warps" must be
+    # sanitized *before* with_knobs() applies it -- with_knobs() does a single
+    # dataclasses.replace() on a frozen, __post_init__-validated config, so an
+    # override that conflicts with the (unrelated, untouched) in_kernel_fc2_reduce
+    # field raises immediately inside with_knobs(), before any post-hoc fixup
+    # here could run.
+    # The knobs dict may itself carry in_kernel_fc2_reduce (e.g. a verbatim
+    # tuned-cache entry); the effective value after with_knobs() is the knob,
+    # not the argument, so sanitize against whichever will win.
+    effective_ikr = (
+        knobs.get("in_kernel_fc2_reduce", in_kernel_fc2_reduce)
+        if knobs
+        else in_kernel_fc2_reduce
+    )
+    if (
+        effective_ikr
+        and knobs
+        and knobs.get("token_back_mode")
+        not in (
+            None,
+            "epi_warps",
+        )
+    ):
+        knobs = {**knobs, "token_back_mode": "epi_warps"}
     cfg = with_knobs(cfg, knobs)
     if cfg.in_kernel_fc2_reduce != in_kernel_fc2_reduce:
-        # Caller-owned correctness choice; see the NVFP4 factory. The MXFP8
-        # kernel rejects ikr together with dispatch-warp token-back, so the
-        # restored ikr also forces epi-warps token-back.
         cfg = dataclasses.replace(
             cfg,
             in_kernel_fc2_reduce=in_kernel_fc2_reduce,
@@ -979,8 +1002,24 @@ def mxfp8_mega_moe(
         raise ValueError(
             f"num_tokens must be in [0, {symm_buffer.num_max_tokens}], got {n}."
         )
-    if n == 0 and symm_buffer._frontend.config.in_kernel_fc2_reduce:
-        return symm_buffer.output_activation[:0] if y is None else None
+    # n == 0 used to shortcut here without ever calling frontend.run() below.
+    # That's unsafe for in_kernel_fc2_reduce: this session's EP peers rely on
+    # every rank physically launching the kernel every round (its persistent
+    # CTA grid -- get_grid_shape() -- is sized from hardware occupancy, not
+    # num_tokens, so even a 0-token round still runs the warp-specialized
+    # dispatch / token-back / tail-cleanup logic peers' cross-rank REDG
+    # combine depends on). A rank that takes this shortcut instead silently
+    # skips that round's participation, desynchronizing the session's
+    # cross-rank bookkeeping -- peers' subsequent launches then wait on a
+    # signal this rank never posts, deadlocking within tens of rounds under
+    # real (unsynchronized, per-rank-independent) traffic.
+    #
+    # n == 0 needs no special case at all: it's just the degenerate instance
+    # of the padding scheme every other n already uses below (stage_inputs()
+    # already fills topk_idx[:capacity] with -1 -- "no work" -- when
+    # num_tokens=0, exactly like it pads topk_idx[n:capacity] for any other
+    # n), so falling through to the same full-buffer frontend.run() call
+    # every nonzero n takes is correct, not just safe.
     if y is not None:
         if y.shape != (n, symm_buffer.hidden):
             raise ValueError(

flashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/nvfp4.py

@@ -1227,8 +1227,27 @@ def nvfp4_mega_moe(
         raise ValueError(
             f"num_tokens must be in [0, {symm_buffer.num_max_tokens}], got {n}."
         )
-    if n == 0 and symm_buffer._frontend.config.fc2_reduces_topk:
-        return symm_buffer.output_activation[:0] if y is None else None
+    # n == 0 used to shortcut here without ever calling frontend.run() below.
+    # That's unsafe for in_kernel_fc2_reduce (fc2_reduces_topk): this
+    # session's EP peers rely on every rank physically launching the kernel
+    # every round (its persistent CTA grid -- get_grid_shape() -- is sized
+    # from hardware occupancy, not num_tokens, so even a 0-token round still
+    # runs the warp-specialized dispatch / token-back / tail-cleanup logic
+    # peers' cross-rank REDG combine depends on -- see the MXFP8 shim's
+    # mxfp8_mega_moe() for the identical bug, root cause, and fix, verified
+    # end-to-end against a real SGLang server). A rank that takes this
+    # shortcut instead silently skips that round's participation,
+    # desynchronizing the session's cross-rank bookkeeping -- peers'
+    # subsequent launches then wait on a signal this rank never posts,
+    # deadlocking within tens of rounds under real (unsynchronized,
+    # per-rank-independent) traffic.
+    #
+    # n == 0 needs no special case at all: it's just the degenerate instance
+    # of the padding scheme every other n already uses below (staging
+    # already marks unrouted rows as "no work" when num_tokens=0, exactly
+    # like it pads the tail for any other n), so falling through to the same
+    # full-buffer frontend.run() call every nonzero n takes is correct, not
+    # just safe.
     if y is not None:
         if y.shape != (n, symm_buffer.hidden):
             raise ValueError(

tests/moe_ep/test_moe_ep_mxfp8_cutedsl_mega_multirank.py

@@ -532,6 +532,207 @@ def test_moe_ep_mxfp8_cutedsl_mega_layer_in_kernel_fc2_reduce():
     )
 
 
+def _run_mega_layer_zero_token_ikr_regression(
+    rank,
+    world_size,
+    *,
+    num_iters: int = 60,
+):
+    """Interleave num_tokens=0 and real forward() calls, in_kernel_fc2_reduce=True,
+    no barrier between iterations, at an independent per-rank schedule.
+
+    Regression guard for mxfp8_mega_moe()'s num_tokens==0 shortcut, which used to
+    return WITHOUT ever calling frontend.run() (i.e. without launching the kernel
+    at all) when in_kernel_fc2_reduce was enabled:
+
+        if n == 0 and symm_buffer._frontend.config.in_kernel_fc2_reduce:
+            return symm_buffer.output_activation[:0] if y is None else None
+
+    Sm100MegaMoEMxfp8Kernel is a persistent megakernel -- its CTA grid
+    (MoEFusedFc12SchedulerParams.get_grid_shape -> (cluster_mn[0], cluster_mn[1],
+    max_active_clusters)) is sized from hardware occupancy, never from
+    num_tokens, so even a genuinely zero-token launch still runs every CTA and
+    the warp-specialized dispatch / token-back / tail-cleanup logic that keeps
+    a rank's cross-rank REDG atomic-add combine session in lockstep with its EP
+    peers. A rank that takes the old shortcut instead silently skips that
+    round's kernel launch, desynchronizing its session state from its peers'
+    -- their subsequent launches then wait on a signal that rank never posts.
+
+    This only manifests when DP/EP ranks call forward() independently (no
+    cross-rank barrier between rounds, exactly how SGLang's per-rank scheduler
+    loop drives it) AND some rank legitimately hits num_tokens==0 (SGLang's own
+    idle-batch mechanism for keeping DP ranks in lockstep) while its peers have
+    real work -- light/symmetric/all-nonzero testing never exercises it. See
+    kernel_src/cutedsl_megamoe/shim/mxfp8.py::mxfp8_mega_moe.
+
+    Shapes/scale intentionally match the real repro (hidden=2048,
+    intermediate=768, num_experts=128, top_k=8, max_tokens_per_rank=16384 --
+    the Qwen3-30B-A3B MXFP8 SGLang config that originally hit this), not this
+    file's usual small test defaults: the same test at
+    hidden=2048/intermediate=1024/num_experts=8/topk=4/max_tokens=64 (this
+    file's ``_mega_problem`` default) plus a deterministic modulo-based
+    zero/nonzero schedule passes vacuously even without the fix -- neither the
+    smaller buffer/expert-count nor a merely-deterministic (as opposed to
+    randomly-timed) schedule reproduces the desync on its own; both the scale
+    and genuinely independent per-rank timing (via per-rank-seeded
+    ``random.Random``, not a fixed formula) were needed to reproduce it.
+
+    CAVEAT: a regression here manifests as a LIVELOCK (100% GPU utilization,
+    zero forward progress, no exception, no crash), not a clean test failure --
+    a same-process watchdog can't interrupt a rank frozen inside
+    torch.cuda.synchronize() on a raw CUDA kernel wait (this isn't an NCCL
+    collective op, so dist.init_process_group(timeout=...) doesn't cover it
+    either). Rely on the CI job's own wall-clock timeout to catch a real
+    regression.
+
+    This exact scenario (matched shapes, matched random schedule, even a
+    deliberate timing nudge on real rounds, NUM_ITERS=2000 -- the script
+    defaults to 60) was confirmed to
+    reliably livelock pre-fix when run as a plain torchrun-launched script
+    (tests/moe_ep/../repro_ikr_zero_token_idle.py -- the authoritative
+    regression artifact for this bug) but passes vacuously pre-fix when run
+    under `torchrun -m pytest` specifically (root cause not fully isolated;
+    pytest's execution environment appears to dampen the wall-clock
+    divergence between ranks this race depends on). This test is kept as a
+    documented, passing correctness check of the exact scenario under the
+    project's normal test harness -- not as a guaranteed regression trap.
+    """
+    import random
+    import time
+
+    import torch
+    import torch.distributed as dist
+
+    from flashinfer.moe_ep import (
+        BootstrapConfig,
+        FleetParams,
+        MegaConfig,
+        MoEEpMegaLayer,
+        MoEEpTensors,
+        MoEWeightPack,
+        ensure_moe_ep_cuda_device,
+    )
+
+    bootstrap = BootstrapConfig(world_size=world_size, rank=rank)
+    ensure_moe_ep_cuda_device(bootstrap)
+
+    hidden = 2048
+    intermediate = 768
+    num_experts = 128
+    topk = 8
+    max_tokens = 16384
+    real_tokens = 4
+    assert num_experts % world_size == 0
+    num_local_experts = num_experts // world_size
+
+    w13, w2 = _make_bf16_weights(
+        rank,
+        num_local_experts=num_local_experts,
+        hidden=hidden,
+        intermediate=intermediate,
+    )
+    warmup_hidden_states, warmup_topk_weights, warmup_topk_ids = _make_inputs(
+        rank, num_tokens=real_tokens, hidden=hidden, num_experts=num_experts, topk=topk
+    )
+    megakernel_config = _megakernel_config(
+        dict(
+            intermediate=intermediate,
+            topk=topk,
+            kind="mxfp8_e4m3",
+            gate_up_clamp=10.0,
+            fast_math=True,
+        ),
+        in_kernel_fc2_reduce=True,
+    )
+
+    mega = MoEEpMegaLayer(
+        bootstrap=bootstrap,
+        fleet_params=FleetParams(
+            num_experts=num_experts,
+            max_tokens_per_rank=max_tokens,
+            token_hidden_size=hidden,
+        ),
+        weights=MoEWeightPack(w13=w13, w2=w2),
+        backend=MegaConfig(megakernel=megakernel_config, preprocess_weights=True),
+    )
+    try:
+        # Matched-count collective warmup -- every rank calls forward() with
+        # real tokens once, together, before the independent-cadence loop.
+        mega.forward(
+            MoEEpTensors(
+                hidden_states=warmup_hidden_states,
+                topk_ids=warmup_topk_ids,
+                topk_weights=warmup_topk_weights,
+            )
+        )
+        torch.cuda.synchronize()
+        dist.barrier()
+
+        # Independently-seeded per-rank RNG, no barrier between iterations:
+        # rank 0 always real (mirrors an always-busy rank); other ranks
+        # independently coin-flip zero/real every iteration, so each rank's
+        # actual wall-clock cadence diverges from its peers' in a way a fixed
+        # formula doesn't produce. Seeded for CI reproducibility.
+        rnd = random.Random(4242 + rank)
+        for it in range(num_iters):
+            n = real_tokens if rank == 0 else (0 if rnd.random() < 0.5 else real_tokens)
+            g = torch.Generator(device="cuda").manual_seed(1000 * it + rank)
+            hidden_states = torch.randn(
+                n, hidden, dtype=torch.bfloat16, device="cuda", generator=g
+            )
+            scores = torch.randn(
+                n, num_experts, dtype=torch.float32, device="cuda", generator=g
+            )
+            topk_weights, topk_ids = torch.topk(
+                scores, topk, dim=-1, largest=True, sorted=False
+            )
+            t = MoEEpTensors(
+                hidden_states=hidden_states,
+                topk_ids=topk_ids.to(torch.int64),
+                topk_weights=topk_weights.to(torch.float32),
+            )
+            if n > 0:
+                # Nudge real per-rank wall-clock divergence: pre-fix, a
+                # zero-token round skips the kernel launch entirely and is
+                # near-instant, while a real round pays actual GPU cost --
+                # under plain torchrun that gap alone is enough to desync
+                # ranks within tens of rounds, but empirically not reliably
+                # under pytest (unconfirmed why; see the CAVEAT above). This
+                # doesn't guarantee detection here, just improves the odds.
+                time.sleep(0.003)
+            y = mega.forward(t)
+            torch.cuda.synchronize()
+            assert y.shape == (n, hidden)
+            assert y.dtype == torch.bfloat16
+            assert torch.isfinite(y).all()
+
+        dist.barrier()
+        return rank
+    finally:
+        mega.destroy()
+
+
+@pytest.mark.gpu_4
+@pytest.mark.arch_blackwell
+def test_moe_ep_mxfp8_cutedsl_mega_layer_in_kernel_fc2_reduce_zero_token_regression():
+    """Zero-token / in_kernel_fc2_reduce livelock regression guard (MXFP8).
+
+    See ``_run_mega_layer_zero_token_ikr_regression`` for the full bug
+    writeup. Before the fix, this reliably livelocks within tens of
+    iterations; after the fix, all ``num_iters`` complete cleanly regardless
+    of each rank's independent zero/nonzero token schedule.
+    """
+    _require_cuda()
+    rank, world_size = _launcher_ranks()
+    if world_size < 4:
+        pytest.skip("needs >=4 ranks")
+    rank = _run_mega_layer_zero_token_ikr_regression(rank, world_size)
+    print(
+        f"rank {rank}: sm100_mxfp8_mxfp8_bf16_cutedsl mega layer survives "
+        "interleaved zero-token/real in_kernel_fc2_reduce forward calls"
+    )
+
+
 @pytest.mark.gpu_4
 @pytest.mark.arch_blackwell
 def test_moe_ep_mxfp8_cutedsl_mega_layer_large_tokens_matches_reference():