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

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.

Start by reviewing #4352 and the gather_qmm_rhs_nax path, especially its conditional tile_matmad_nax behavior for inactive simdgroups. Compare v0.32.2 with the guarded revert at ceab91938, then run the package/kernel/import tests and the sustained heterogeneous M3 Ultra + M5 Max JACCL workload. Done means the 30K follow-up and other listed collective rotations complete without a lost completion.