[GLM] Rework MoE routing to avoid scatter_/gather lowerings on multi-row mesh

[mvasiljevicTT]

Ticket

Fixes #5409

Problem description

GLM 4.7 had good PCC (0.99) only on the first 16/64 users — those on the first of the 4 devices the batch is sharded across — because the MoE router indexes a global tensor with a large index that is either mis-sharded or precision-truncated. Two bugs of this class, both in route_tokens_to_experts (sparse_mlp.py):

1. group_mask.scatter_(1, group_idx, 1) — scatter_ adds a token row-iota that Shardy shards without the per-shard offset, so the group mask marks only mesh-row 0 and zeros routing for rows 1-3 (0.88 → 0.95).
2. router_logits.gather(1, topk_indices) — lowers to ttir.embedding over the all-gathered [tokens·E] score table with a flat index token·E + expert computed at fp16-class precision. For rows 1-3 the large token·E offset rounds off the expert bits (e.g. 2609 → 2610), gathering the wrong expert's score (0.95 → 0.99).

What's changed

Replaces both ops with the one_hot pattern already used for XLA compat (_topk_to_sparse_scores), so the arange is over experts/groups (a model dim, not the sharded batch axis) and the index stays small + exact:

• Group mask: scatter_ → one_hot + any.
• Top-k weights: gather → one_hot + einsum.

Both rewrites are mathematically equivalent on CPU and just workarounds around buggy scatter and gather ops to unblock glm pcc until they are fixed.

Impact

• GLM 4.7 multi-device PCC 0.88 → 0.993 on all 4 mesh-rows (full-batch decode PCC 0.9937); benchmark PASSES.
• CI A/B (GLM perf, shared runner off): no fix → fails rows 1-3 · with fix → passes

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 33.80%. Comparing base (1bd6e3b) to head (48adda3).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #5435   +/-   ##
=======================================
  Coverage   33.80%   33.80%           
=======================================
  Files          37       37           
  Lines        4990     4990           
=======================================
  Hits         1687     1687           
  Misses       3303     3303           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.