qwen3_5_moe_split: drain — eval gathered instead of contribution

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
dmcc73andClaude Opus 4.7 committed 2026-04-21 17:28:19 +01:00
1 parent 95d4e66680
commit 312229a454
1 file changed
+1 -1
@@ -305,8 +305,8 @@ def pipelined_layer_loop(
contribution = moe(last_layer, h_H1_pending)
else:
contribution = _zeros_like(h_H1_pending)
mx.eval(contribution)
gathered = mx.distributed.all_gather(contribution, group=group)
mx.eval(gathered)
out_H1 = gathered[MOE_RANK : MOE_RANK + 1]
# After the final B stage (Stage 2N-1), out_{N-1}_H0 = x_H0 (MOE's contribution