Moe Permute Batch Scaling
Keep native MoE permutation latency from scaling sharply with batch size.
I'm investigating a performance issue with the native _moe_C.moe_permute operator on our A100 workload. The aligned-routing path produces the correct results, but its latency grows disproportionately as I increase the token batch. Could you find what's causing this and remove the scaling bottleneck?
I need the operator to keep supporting the same inputs and preserve its observable behavior in both aligned and unaligned cases, including exact expert offsets, inverse and permuted mappings, expert ranges, payload bytes, and sentinels. Please don't introduce a new expert-count limit. The algorithm, helper names, and kernel decomposition are up to you.
The checkout is in /workspace/vllm. Please rebuild the focused native _moe_C extension using the repository's CUDA, CUTLASS, and CMake build flow, then check the rebuilt candidate's correctness before timing it. Include its SHA-256 and the path it loads from in a fresh process so I can identify the binary you measured.
For the performance check, use one NVIDIA A100-SXM4-40GB with 64 experts, top-k 6, hidden width 2048, FP8 E4M3 payload storage, and expert ranges aligned to 128 rows. Distribute routes cyclically; for example, token t and slot k can route to (17*t + 7*k) mod 64. Measure token counts 1, 32, 128, 512, 1024, 2048, 4096 with 20 warmups and the median of five trials of 50 iterations. I need the 4096-token median below 250 us and the 4096/512 latency ratio below 3.5. These are the shapes I'm using to measure performance; correctness still needs to hold for the operator's other supported inputs.