2102146c01
thrust/examples/mode.cu entire design (80 lines):
Complete GPU pipeline: sort → unique_count → reduce_by_key → max_element
Key operations:
1. thrust::sort — bring equal keys together (we already do: argsort)
2. thrust::unique_count — precompute number of unique keys for allocation
3. thrust::reduce_by_key(data, constant_iterator<1>) — count per key
4. thrust::max_element — find the mode (highest count)
Design principle: every step is a GPU primitive, no CPU round-trips.
constant_iterator<1> trick: turns reduce_by_key into count_by_key.
Translation to MoE segment detection:
Previous (4 GPU ops + CPU tensors):
changes = cat([True, sorted[1:] != sorted[:-1]])
seg_starts = changes.nonzero()
seg_ends = cat([seg_starts[1:], tensor([len])])
seg_eids = sorted[seg_starts]
Now (1 fused GPU op):
seg_eids, _, seg_counts = torch.unique_consecutive(sorted, return_counts=True)
seg_ends = seg_counts.cumsum(0)
seg_starts = cat([0, seg_ends[:-1]])
unique_consecutive IS mode.cu's sort+reduce_by_key fused: it returns
(unique_keys, inverse, counts) — exactly the data mode.cu builds from
reduce_by_key(data, constant_iterator<1>, keys_out, counts_out).
3 fewer GPU kernel launches per MoE forward.
CCCL source: thrust/examples/mode.cu
Maps to: qwen3_6_scripts/qwen3_5.py (_pure_pytorch_experts)