Packed Sort: speed up MoE routing argsort via bit-packed single-key sort

[bigeye]

MoE routing sorts tokens by expert with a stable argsort, which XLA runs
as a two-operand (key, index) sort. When the key and the index fit
together in an int32, pack them into one value, key << shift | index,
and run a plain single-key unstable sort instead. The index in the low
bits breaks ties, so the result is identical to the stable argsort, but
XLA has no iota operand and no stable-sort comparator to deal with.

The argsort for gather drops from 106.6 us to 21.7 us (4.9x speedup),
and the validity sort from 7.2us to 5.5us (1.3x speedup).

Measured on Qwen/Qwen3.5-397B-A17B-FP8, TPU v7x-8, preset 8k1k, 640
prompts, repeat=5, The two sides ran with:

• Baseline: without any change. no env overrides.
• Packed Sort + FP8 AG: MOE_ALL_GATHER_ACTIVATION_DTYPE=fp8
  SC_ALLREDUCE_ALLGATHER_OFFLOAD_MIN_BYTES=0

MoE AG and Sort/Permute runs in parallel, so we need to optimize both at
the same time. Read the e2e deltas as the two changes together, not as the
sort alone.

Conc	tok/s	std%	TTFT (ms)	TPOT (ms)
64	15908.5 -> 16114.74 (+1.30%)	0.99 -> 0.98	2014.13 -> 1999.67 (-0.72%)	31.305 -> 30.906 (-1.27%)
128	22883.2 -> 23027.63 (+0.63%)	0.53 -> 0.56	2242.35 -> 2114.58 (-5.70%)	44.922 -> 44.642 (-0.62%)
256	28007.1 -> 30381.17 (+8.48%)	5.48 -> 0.87	3258.17 -> 3127.76 (-4.00%)	65.219 -> 64.296 (-1.42%)
512	30888.1 -> 31775.83 (+2.87%)	0.34 -> 0.39	30172.15 -> 28997.39 (-3.89%)	94.396 -> 90.833 (-3.77%)

Raw per-repeat throughputs for Packed Sort + FP8 AG:
c64 [15860.9, 16083.5, 16238.5, 16247.4, 16143.4]
c128 [22798.9, 23088.1, 23098.6, 23052.7, 23099.9]
c256 [30270.5, 30141.1, 30672.1, 30660.8, 30161.4]
c512 [31665.6, 31892.7, 31779.7, 31902.3, 31638.9]

Eval	Baseline	Packed Sort + FP8 AG
mmlu_pro (700 q) *	82.8	82.8
gpqa_diamond_cot_zeroshot (50 q, flexible-extract)	88.0 +/- 1.6	92.0 +/- 1.6
aime25 (30 q)	73.3 +/- 8.6	77.5 +/- 4.2

Signed-off-by: Donghyun Cho donghyun@google.com

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

If the change fixes a Github issue, please include a link, e.g.,:
FIXES: #123456

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.

[bigeye]

@qizzzh @guowei-dev for reviews.

[guowei-dev]

We've been exploring similar packed-sort ideas on the torchtpu side — https://github.com/vllm-project/vllm-torchtpu/pull/649 does the same routing packing, plus more of the routing prep, including a counting sort for the inverse permutation. If this is headed for torchtpu as well, maybe we can build on that directly?

[bigeye]

@guowei-dev Thanks for the heads up! Good to know that there's some similar efforts being proven. I actually have a branch (will have a separate pr for that) for counting sort for the second sort as well. I've been working on that and saw the good performance gain there. As they are different repos at the end of the day, I think we can just apply these techniques in both repo in that it won't hurt the migration or future plans. Tell me if you have other thoughts.

[guowei-dev]

LGTM per offline discussion. The changes itself are in good shape, thanks for working on this!