[Kernel] Add prepared-input fast path with MiniMax-M2 top-k/act quant fusion

[qianlihuang]

Purpose

This PR adds a MiniMax-M2 decode fast path for the supported modular FP8 MoE path.

There are two parts to the change:

1. Adds a narrow prepared-input hook to shared FusedMoE, similar in spirit to DeepSeek V4 MegaMoE’s prepare_megamoe_inputs. MiniMax-M2 uses shared FusedMoE, so this avoids building a separate model-specific MoE path just for prepared routing/quant inputs. The same hook can also be reused by other shared-FusedMoE models with similar router/FP8 staging if they add their own staging kernels later.

2. Fuses two independent small-batch decode operations: sigmoid+bias top-k routing and per-token-group FP8 activation quantization. The fused Triton op produces prepared top-k weights/ids plus FP8 activations/scales, then the MoE runner consumes them directly. The optimization is gated to the MiniMax-M2 Triton FP8 MoE path.

bs=16, TP4 H800

Before ~8us
dp0_pp0_tp0_dcp0_ep0_rank0.1779861662605309419.pt.trace.json.gz

After ~3us
dp0_pp0_tp0_dcp0_ep0_rank0.1779875418507998155.pt.trace.json.gz

Test Plan

microbench

python3 benchmarks/kernels/benchmark_minimax_moe_topk_quant.py \
  --m-values 1 2 4 8 16 32 64 128 256 512 1024 \
  --modes eager cudagraph

serve

VLLM_USE_DEEP_GEMM_E8M0=0 vllm serve MiniMaxAI/MiniMax-M2.7   -tp 4   --tool-call-parser minimax_m2   --reasoning-parser minimax_m2   --enable-auto-tool-choice   --trust-remote-code   --no-enable-prefix-caching

bench

python3 tests/evals/gsm8k/gsm8k_eval.py 

vllm bench serve --backend openai-chat --model "$MODEL" --endpoint /v1/chat/completions --host 127.0.0.1 --port "$PORT" --dataset-name random --num-prompts 128 --random-input-len 8192 --random-output-len 1024 --random-range-ratio 0.0 --num-warmups 16 --request-rate inf --max-concurrency 16 --save-result

Test Result

mode            M      H     E   K  blk    topk_us   quant_us    base_us   fused_us   speedup   delta_us   id_mis      w_err      q_err      s_err
---------------------------------------------------------------------------------------------------------------------------------------------------
eager           1   3072   256   8  128      8.928      6.272     12.256      6.912     1.773      5.344        0   1.49e-08   0.00e+00   0.00e+00
eager           2   3072   256   8  128      9.088      6.112     12.480      6.880     1.814      5.600        0   1.49e-08   0.00e+00   0.00e+00
eager           4   3072   256   8  128      9.568      6.144     13.120      7.232     1.814      5.888        0   1.49e-08   0.00e+00   0.00e+00
eager           8   3072   256   8  128     10.272      6.240     13.504      7.008     1.927      6.496        0   1.49e-08   0.00e+00   0.00e+00
eager          16   3072   256   8  128     10.112      6.528     13.472      7.360     1.830      6.112        0   1.49e-08   0.00e+00   0.00e+00
eager          32   3072   256   8  128     10.528      6.656     13.952      7.520     1.855      6.432        0   2.98e-08   0.00e+00   0.00e+00
eager          64   3072   256   8  128     10.240      6.496     14.080      7.360     1.913      6.720        0   4.47e-08   0.00e+00   0.00e+00
eager         128   3072   256   8  128     10.432      7.008     14.528      8.864     1.639      5.664        0   4.47e-08   0.00e+00   0.00e+00
eager         256   3072   256   8  128     11.264      7.360     15.712     10.912     1.440      4.800        0   4.47e-08   0.00e+00   0.00e+00
eager         512   3072   256   8  128     11.360      9.312     17.856     15.104     1.182      2.752        0   4.47e-08   0.00e+00   0.00e+00
eager        1024   3072   256   8  128     12.832     12.064     21.312     23.712     0.899     -2.400        0   5.96e-08   0.00e+00   0.00e+00
cudagraph       1   3072   256   8  128      4.653      1.735      6.389      2.828     2.260      3.562        0   1.49e-08   0.00e+00   0.00e+00
cudagraph       2   3072   256   8  128      4.659      1.821      6.494      2.833     2.292      3.661        0   1.49e-08   0.00e+00   0.00e+00
cudagraph       4   3072   256   8  128      5.250      1.840      6.955      2.857     2.434      4.098        0   1.49e-08   0.00e+00   0.00e+00
cudagraph       8   3072   256   8  128      5.532      1.857      7.466      2.869     2.602      4.597        0   1.49e-08   0.00e+00   0.00e+00
cudagraph      16   3072   256   8  128      5.668      1.881      7.507      2.910     2.580      4.597        0   1.49e-08   0.00e+00   0.00e+00
cudagraph      32   3072   256   8  128      5.868      1.956      7.818      2.960     2.641      4.858        0   2.98e-08   0.00e+00   0.00e+00
cudagraph      64   3072   256   8  128      5.925      2.016      7.958      3.069     2.593      4.888        0   4.47e-08   0.00e+00   0.00e+00
cudagraph     128   3072   256   8  128      6.035      2.152      8.196      4.432     1.849      3.765        0   4.47e-08   0.00e+00   0.00e+00
cudagraph     256   3072   256   8  128      6.180      2.447      8.596      6.311     1.362      2.285        0   4.47e-08   0.00e+00   0.00e+00
cudagraph     512   3072   256   8  128      6.661      3.688     10.232     10.187     1.004      0.045        0   4.47e-08   0.00e+00   0.00e+00
cudagraph    1024   3072   256   8  128      7.830      5.484     13.400     18.574     0.721     -5.174        0   5.96e-08   0.00e+00   0.00e+00

Metric	Before	After
gsm8k	~0.92	0.921

Before

============ Serving Benchmark Result ============
Successful requests:                     128       
Failed requests:                         0         
Maximum request concurrency:             16        
Benchmark duration (s):                  150.75    
Total input tokens:                      1053440   
Total generated tokens:                  131072    
Request throughput (req/s):              0.85      
Output token throughput (tok/s):         869.46    
Peak output token throughput (tok/s):    1056.00   
Peak concurrent requests:                29.00     
Total token throughput (tok/s):          7857.44   
---------------Time to First Token----------------
Mean TTFT (ms):                          1808.00   
Median TTFT (ms):                        1903.28   
P99 TTFT (ms):                           4204.88   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          16.64     
Median TPOT (ms):                        16.62     
P99 TPOT (ms):                           18.17     
---------------Inter-token Latency----------------
Mean ITL (ms):                           21.97     
Median ITL (ms):                         15.67     
P99 ITL (ms):                            233.70    
==================================================

After

============ Serving Benchmark Result ============
Successful requests:                     128       
Failed requests:                         0         
Maximum request concurrency:             16        
Benchmark duration (s):                  146.46    
Total input tokens:                      1053440   
Total generated tokens:                  131072    
Request throughput (req/s):              0.87      
Output token throughput (tok/s):         894.93    
Peak output token throughput (tok/s):    1088.00   
Peak concurrent requests:                25.00     
Total token throughput (tok/s):          8087.57   
---------------Time to First Token----------------
Mean TTFT (ms):                          1571.37   
Median TTFT (ms):                        1663.23   
P99 TTFT (ms):                           4207.63   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          16.35     
Median TPOT (ms):                        16.25     
P99 TPOT (ms):                           17.62     
---------------Inter-token Latency----------------
Mean ITL (ms):                           22.21     
Median ITL (ms):                         15.41     
P99 ITL (ms):                            236.05    
==================================================

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[gemini-code-assist]

Code Review

This pull request introduces optimizations for the MiniMax-M2 model by enabling the pre-computation of MoE inputs, including top-k weights, IDs, and quantized activations, which are then passed directly to fused MoE kernels. Key changes include a new CUDA kernel for expert alignment in small-batch scenarios, a specialized Triton kernel for fused top-k sigmoid and FP8 quantization, and significant updates to the MoE runner and modular kernel infrastructure to support these prepared inputs. Feedback focuses on ensuring the Triton kernel robustly handles cases where hidden size is not a multiple of the block size and addressing potential silent performance fallbacks if input padding is inconsistent between the model preparation logic and the MoE runner.

[Copilot]

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Adds a specialized fused top‑k (sigmoid + bias) routing + FP8 activation quantization path for MiniMax‑M2, and threads “prepared” MoE inputs through the fused MoE stack to reuse precomputed routing/quantization.

Changes:

• Introduces a new Triton custom op that fuses top‑k routing and per-token-group FP8 quantization, with a dispatch threshold.
• Adds “prepared_*” plumbing across MoE runner / kernel / quant methods so callers can skip internal routing/quantization when precomputed.
• Adds a standalone benchmark for the new fused kernel.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
vllm/model_executor/models/minimax_m2.py	Detects when MiniMax‑M2 can use prepared MoE inputs and calls the fused topk+quant dispatch op.
vllm/model_executor/layers/quantization/fp8.py	Exposes supports_prepared_inputs and forwards prepared tensors into the MoE kernel apply path.
vllm/model_executor/layers/fused_moe/topk_quant_kernels.py	New Triton kernel + custom ops to fuse top‑k routing and FP8 quantization (plus baseline fallback).
vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py	Extends runner interface to accept prepared routing/quant inputs.
vllm/model_executor/layers/fused_moe/runner/moe_runner.py	Validates/discards incompatible prepared inputs and threads prepared tensors through runner execution.
vllm/model_executor/layers/fused_moe/router/base_router.py	Adds post-processing helper for externally-computed top‑k results.
vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py	Implements prepared-input preparation for the no-dp-ep path.
vllm/model_executor/layers/fused_moe/modular_kernel.py	Adds prepared-input support plumbing into the modular MoE kernel, including capability checks.
vllm/model_executor/layers/fused_moe/layer.py	Forwards prepared inputs from layer forward into the runner.
vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py	Exposes supports_prepared_inputs and forwards prepared tensors to the MoE kernel.
vllm/model_executor/layers/fused_moe/fused_moe_method_base.py	Adds base API surface for prepared inputs.
benchmarks/kernels/benchmark_minimax_moe_topk_quant.py	Benchmark harness comparing fused vs baseline routing+quantization.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[yewentao256]

Thanks for the work!

[yewentao256]

use_ue8m0=False Could you also test with use_ue8m0=True?

[qianlihuang]

Thanks for the suggestion. MiniMax-M2.7 uses the non-UE8M0 FP8 MoE path, and the fused kernel in this PR intentionally matches per_token_group_quant_fp8(..., use_ue8m0=False). I also benchmarked standalone per_token_group_quant_fp8 with use_ue8m0=True vs False on the MiniMax-M2 shape (M, 3072), block_k=128. The latency is very close.

Command and results

PYTHONPATH=. python3 - <<'PY'
import torch
from vllm.model_executor.layers.quantization.utils.fp8_utils import per_token_group_quant_fp8
from vllm.triton_utils import triton

ms = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024]
H = 3072
block_k = 128
N_RUNS = 3

def bench_us(fn):
    results = []
    for _ in range(N_RUNS):
        r = triton.testing.do_bench(fn)
        if isinstance(r, tuple):
            r = r[0]
        results.append(float(r) * 1000.0)
    return sum(results) / len(results)

print(f"{'M':>6} {'ue8m0=false us':>16} {'ue8m0=true us':>16} {'ratio':>8}")
for m in ms:
    x = torch.randn((m, H), device="cuda", dtype=torch.bfloat16)
    torch.cuda.synchronize()

    t_false = bench_us(
        lambda: per_token_group_quant_fp8(x, block_k, use_ue8m0=False)
    )
    t_true = bench_us(
        lambda: per_token_group_quant_fp8(x, block_k, use_ue8m0=True)
    )

    print(f"{m:6d} {t_false:16.3f} {t_true:16.3f} {t_true / t_false:8.3f}")
PY

     M   ue8m0=false us    ue8m0=true us    ratio
     1            6.024            6.226    1.034
     2            6.262            6.234    0.995
     4            6.304            6.271    0.995
     8            6.149            6.542    1.064
    16            6.302            6.385    1.013
    32            6.443            6.705    1.041
    64            6.744            6.616    0.981
   128            6.790            7.102    1.046
   256            7.421            7.707    1.038
   512            9.638            9.785    1.015
  1024           11.990           11.862    0.989

[yewentao256]

Would this PR changes the behavior of e8m0 = True path? Please take a check

[qianlihuang]

Yes, this can change the E8M0 path since the fused kernel matches use_ue8m0=False. And I think False is more appropriate for MiniMax-M2.7, but I can add a guard to keep E8M0 unchanged if preferred.

[yewentao256]

As E8M0 is widely used currently, I am not sure if we really want this complicated optimization for non-e8m0 path. Also CC @mgoin

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @qianlihuang.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork