7ad59e781f
Qwen3.6-35B-A3B has 256 experts × top_k=8. The baseline prefill MoE:
for eid in unique_eids: # up to 256 iterations
tokens = hidden_states[tok_ids] # SCATTERED gather
F.linear(tokens, w13[eid])
Problem: hidden_states[tok_ids] creates a non-contiguous gather for each expert.
With 16384 tokens × 256 experts, this is 256 scattered gathers per layer.
Optimization (CCCL segmented-sort pattern):
1. Flatten all token-expert pairs: (T×K,) assignments
2. Sort by expert ID: tokens for same expert become CONTIGUOUS
3. Each F.linear gets contiguous input → much better memory access
4. Activation (silu × up) computed in ONE fused op across all pairs
5. index_add_ scatter-back is one kernel call
Memory access improvement:
Before: 256 × hidden_states[random_indices] → scattered HBM reads
After: sorted_tokens[start:end] → sequential HBM reads per expert
The expert loop still exists (can't batch variable-size GEMMs with F.linear),
but each iteration reads contiguous memory instead of scattered indices.