When ici_context_parallelism > 1, the sequence axis is sharded across devices. RoutedMoE flattens (batch, sequence) into a single token axis before the grouped matmul. The sharding on the sequence axis doesn't survive that reshape, so XLA gathers the full sequence onto each device.
The result is correct but the memory saving from context parallelism is lost in the MoE block.
Where
Qwen3NextSparseMoeBlock calls RoutedMoE for the routed experts. Any model that combines RoutedMoE with context parallelism is affected.
What I did instead
I substituted a plain MlpBlock for RoutedMoE when context parallelism is active. This preserves the sequence sharding and let a long-context run proceed.
Don't take this as the fix. It's slower. Measured on a v5p at sequence 8192, the grouped matmul hits 29.01% MFU and MlpBlock hits 27.82%, so the swap costs real throughput. It also only works for a dense config, which has nothing to route.
The real fix is inside RoutedMoE: keep the sequence sharding through the grouped matmul. That needs whoever owns that kernel.
Reproduce
Run any model that uses RoutedMoE with ici_context_parallelism=4 and a sequence that doesn't fit on one device. The step runs out of HBM. Raising ici_context_parallelism saves no memory.
cc @mmcsa
When
ici_context_parallelism > 1, the sequence axis is sharded across devices.RoutedMoEflattens(batch, sequence)into a single token axis before the grouped matmul. The sharding on the sequence axis doesn't survive that reshape, so XLA gathers the full sequence onto each device.The result is correct but the memory saving from context parallelism is lost in the MoE block.
Where
Qwen3NextSparseMoeBlockcallsRoutedMoEfor the routed experts. Any model that combinesRoutedMoEwith context parallelism is affected.What I did instead
I substituted a plain
MlpBlockforRoutedMoEwhen context parallelism is active. This preserves the sequence sharding and let a long-context run proceed.Don't take this as the fix. It's slower. Measured on a v5p at sequence 8192, the grouped matmul hits 29.01% MFU and
MlpBlockhits 27.82%, so the swap costs real throughput. It also only works for a dense config, which has nothing to route.The real fix is inside
RoutedMoE: keep the sequence sharding through the grouped matmul. That needs whoever owns that kernel.Reproduce
Run any model that uses
RoutedMoEwithici_context_parallelism=4and a sequence that doesn't fit on one device. The step runs out of HBM. Raisingici_context_parallelismsaves no memory.cc @mmcsa