[BUG] NAX MoE simdgroup skipping in 0.32.2 triggers lost JACCL completion under heterogeneous TP

[jonathan308]

Describe the bug

MLX v0.32.2 reproducibly loses a JACCL completion during sustained tensor-parallel MoE inference on a heterogeneous M3 Ultra + M5 Max pair. An exact A/B isolates the regression to #4352 (c7ff35d97), which skips tile_matmad_nax for inactive simdgroups in gather_qmm_rhs_nax.

The M3 rank does not use NAX; the M5 rank does. Reverting only #4352 on top of v0.32.2 makes the full failure sequence pass while retaining every other 0.32.2 change.

Hardware and software

• rank 0: Apple M3 Ultra, 256 GiB
• rank 1: Apple M5 Max, 128 GiB
• macOS 26.5.2 on both
• direct Thunderbolt RDMA, JACCL, two equal tensor-parallel ranks
• model: DeepSeek-V4-Flash MXFP4 MoE, 43 layers, about 85.6 GB resident per rank
• MLX-LM 0.31.3
• MTP speculative decode depth 5

The MLX builds also include a progress timeout and completion-status diagnostics. The timeout turns the otherwise indefinite spin into exit 75 after 30 seconds without a completion; no error work completion was reported before the timeout.

Exact A/B

v0.32.2 with #4352 present

• package/kernel/import tests pass
• a 120,008-token cold prefill can complete at about 651 tok/s
• the following cold 30K request loses an all_reduce completion and returns zero output tokens
• after rebooting the M5, reloading from clean state, and issuing a cold 30K request, the same failure reproduces
• rank 1 exits after all_reduce made no progress for 30001ms
• about 92–94 GiB remains wired on the M5 until reboot

Result: 2/2 physical 30K failure reproductions.

v0.32.2 with only #4352 reverted

Commit used for the tested guarded build: ceab91938.

• 3,440 dependent JACCL collectives: correct on both ranks, 0.829 seconds
• cold 30K: 736.3 prefill tok/s, 72.0 API decode tok/s, 64/64 tokens
• cold 120K: 659.8 prefill tok/s, 59.6 API decode tok/s, 64/64 tokens
• immediately following cold 30K: completed, 73.5 API decode tok/s
• third cold 30K rotation: 717.7 prefill tok/s, 72.8 API decode tok/s
• exact 30K prefix reuse: 30,012 cached tokens, 0.82 second TTFT
• two concurrent 128-token decoders: 72.8 aggregate tok/s, both correct
• decode plus 24K cold prefill: decoder 77.6 tok/s; prefill 676.1 tok/s; both correct
• zero failed requests and no lost completion

Result: the original 120K -> 30K failure sequence passes, along with cache and concurrency rotations.

Suspected boundary

#4352 conditionally omits NAX matmul instructions for inactive simdgroups inside a threadgroup. The published single-machine M5 benchmarks pass, but the sustained heterogeneous TP workload appears to leave the M5 GPU/RDMA sequence unable to complete a later collective. The non-NAX M3 rank never takes that path.

This is an inference from the exact single-commit revert A/B; I do not yet have a smaller standalone kernel reproducer.

Expected behavior

The optimized MoE kernel must not poison later Metal/JACCL work. Until the inactive-simdgroup optimization can be made safe under this sustained workload, the pre-#4352 implementation is the known-good path.

I have a minimal revert ready and can help test a narrower kernel fix or runtime-gated variant on the physical M3 Ultra + M5 Max RDMA pair.

[nastya236]

Thank you for reporting the bug.
Skipping simdgroups substantially improved the prefill and decode for quantized MoEs on m5. Therefore we will keep this change. We will think about and implement fallbacks for heterogeneous set ups.