[MoE][Kernel] Add optional HPC BF16xFP32 router GEMM

[lengrongfu]

Purpose

Issues: #49277

Add an opt-in HPC BF16xFP32 GEMM path for FP32 MoE router weights on SM90 GPUs.

Guard the path by HPC availability, device capability, dtype, layout, and shape requirements. Cache the BF16 high/low weight
decomposition and preserve the existing router GEMM fallback when the HPC kernel is unavailable or ineligible.

Reproduce

1. install hpc-ops

git clone https://github.com/Tencent/hpc-ops.git
cd hpc-ops

# build packages
make wheel
python3 -m pip install dist/*.whl

2. start vllm serve

export VLLM_ENABLE_HPC_ROUTER_GEMM=1
$ vllm serve --model /new-model/MiniMax-M3-MXFP8/ --trust-remote-code --block-size 128 --tensor-parallel-size 8 --tool-call-parser minimax_m3 --enable-auto-tool-choice --reasoning-parser minimax_m3 --max-model-len 65535

3. vllm bench

vllm bench serve --served-model-name /new-model/MiniMax-M3-MXFP8/ --model /new-model/MiniMax-M3-MXFP8/ --tokenizer /new-model/MiniMax-M3-MXFP8/ --backend openai-chat --endpoint /v1/chat/completions --dataset-name random --random-input 10000 --random-output 1 --seed 131 --base-url http://localhost:8000 --num-prompts 100

Performance

Model: MiniMax-M3-MXFP8
Workload: 100 concurrent requests

Prefill-heavy workload

Input:  10000 tokens
Output: 1 token

Metric	Original	HPC	Change
Total throughput	8471.53 tok/s	8667.48 tok/s	+2.31%
Mean TTFT	61938.32 ms	60540.02 ms	+2.26%
Median TTFT	62196.86 ms	60783.03 ms	+2.27%
P99 TTFT	119626.03 ms	116913.21 ms	+2.27%

Decode-heavy workload

Input:  1 token
Output: 100 tokens

Metric	Original	HPC	Change
Mean TPOT	64.03 ms	63.04 ms	+1.55%
Median TPOT	63.97 ms	62.99 ms	+1.53%
P99 TPOT	64.03 ms	63.05 ms	+1.53%
Mean ITL	63.83 ms	62.76 ms	+1.68%
Total throughput	3251.03 tok/s	3340.22 tok/s	+2.74%

Model Evaluation

GSM8K 5-shot evaluation was run on all 1,319 questions using
MiniMax-M3-MXFP8.

Metric	Original	HPC	Change
Accuracy	0.879	0.882	+0.003
Invalid responses	0.000	0.000	No change
Total latency	121.214 s	118.726 s	-2.05%
Questions per second	10.882	11.110	+2.10%
Output tokens per second	978.483	997.535	+1.95%

No accuracy regression was observed. The 0.3 percentage-point difference is small relative to the expected run-to-run variation and is not treated as a statistically significant accuracy improvement.

The HPC run improved evaluation throughput by approximately 2%.

lm_eval

$ lm_eval   --model local-completions   --model_args "pretrained=/new-model/MiniMax-M3-MXFP8/,tokenizer=/new-model/MiniMax-M3-MXFP8/,base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=50,max_retries=3"   --tasks gsm8k   --num_fewshot 5   --apply_chat_template   --fewshot_as_multiturn   --seed 0

• disable hpc-ops

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9515	±	0.0059
		strict-match	5	exact_match	↑	0.9507	±	0.0060

• enable hpc-ops

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9522	±	0.0059
		strict-match	5	exact_match	↑	0.9522	±	0.0059

Test Plan

• Prefill test

vllm bench serve --served-model-name /new-model/MiniMax-M3-MXFP8/ --model /new-model/MiniMax-M3-MXFP8/ --tokenizer /new-model/MiniMax-M3-MXFP8/ --backend openai-chat --endpoint /v1/chat/completions --dataset-name random --random-input 10000 --random-output 1 --seed 131 --base-url http://localhost:8000 --num-prompts 100

• Decode test

vllm bench serve --served-model-name /new-model/MiniMax-M3-MXFP8/ --model /new-model/MiniMax-M3-MXFP8/ --tokenizer /new-model/MiniMax-M3-MXFP8/ --backend openai-chat --endpoint /v1/chat/completions --dataset-name random --random-input 1 --random-output 100 --seed 132 --base-url http://localhost:8000 --num-prompts 100

3 * Benchmark

• benchmark command

vllm bench serve --served-model-name /new-model/MiniMax-M3-MXFP8/ --model /new-model/MiniMax-M3-MXFP8/ --tokenizer /new-model/MiniMax-M3-MXFP8/ --backend openai-chat --endpoint /v1/chat/completions --dataset-name random --random-input 1 --random-output 100 --seed 132 --base-url http://localhost:8000 --num-prompts 100

disable hpc-ops

• one

============ Serving Benchmark Result ============
Successful requests:                     100       
Failed requests:                         0         
Benchmark duration (s):                  13.37     
Total input tokens:                      117600    
Total generated tokens:                  1000      
Request throughput (req/s):              7.48      
Output token throughput (tok/s):         74.78     
Peak output token throughput (tok/s):    371.00    
Peak concurrent requests:                100.00    
Total token throughput (tok/s):          8868.88   
---------------Time to First Token----------------
Mean TTFT (ms):                          7082.05   
Median TTFT (ms):                        7244.71   
P99 TTFT (ms):                           12051.08  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          555.21    
Median TPOT (ms):                        547.82    
P99 TPOT (ms):                           836.48    
---------------Inter-token Latency----------------
Mean ITL (ms):                           499.69    
Median ITL (ms):                         627.93    
P99 ITL (ms):                            846.17    
==================================================

• two

============ Serving Benchmark Result ============
Successful requests:                     100       
Failed requests:                         0         
Benchmark duration (s):                  12.14     
Total input tokens:                      117600    
Total generated tokens:                  1000      
Request throughput (req/s):              8.24      
Output token throughput (tok/s):         82.40     
Peak output token throughput (tok/s):    394.00    
Peak concurrent requests:                100.00    
Total token throughput (tok/s):          9772.63   
---------------Time to First Token----------------
Mean TTFT (ms):                          5904.29   
Median TTFT (ms):                        6021.40   
P99 TTFT (ms):                           10822.47  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          549.59    
Median TPOT (ms):                        545.39    
P99 TPOT (ms):                           830.11    
---------------Inter-token Latency----------------
Mean ITL (ms):                           494.63    
Median ITL (ms):                         641.32    
P99 ITL (ms):                            837.92    
==================================================

• three

============ Serving Benchmark Result ============
Successful requests:                     100       
Failed requests:                         0         
Benchmark duration (s):                  12.11     
Total input tokens:                      117600    
Total generated tokens:                  1000      
Request throughput (req/s):              8.25      
Output token throughput (tok/s):         82.55     
Peak output token throughput (tok/s):    394.00    
Peak concurrent requests:                100.00    
Total token throughput (tok/s):          9790.00   
---------------Time to First Token----------------
Mean TTFT (ms):                          5895.77   
Median TTFT (ms):                        6019.90   
P99 TTFT (ms):                           10804.74  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          548.48    
Median TPOT (ms):                        543.73    
P99 TPOT (ms):                           827.84    
---------------Inter-token Latency----------------
Mean ITL (ms):                           493.63    
Median ITL (ms):                         645.86    
P99 ITL (ms):                            831.84    
==================================================

enable hpc-ops

• one

============ Serving Benchmark Result ============
Successful requests:                     100       
Failed requests:                         0         
Benchmark duration (s):                  11.87     
Total input tokens:                      117600    
Total generated tokens:                  1000      
Request throughput (req/s):              8.43      
Output token throughput (tok/s):         84.25     
Peak output token throughput (tok/s):    313.00    
Peak concurrent requests:                100.00    
Total token throughput (tok/s):          9992.63   
---------------Time to First Token----------------
Mean TTFT (ms):                          5795.40   
Median TTFT (ms):                        5846.14   
P99 TTFT (ms):                           10573.04  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          534.39    
Median TPOT (ms):                        536.62    
P99 TPOT (ms):                           807.26    
---------------Inter-token Latency----------------
Mean ITL (ms):                           480.95    
Median ITL (ms):                         688.79    
P99 ITL (ms):                            810.45    
==================================================

• two

============ Serving Benchmark Result ============
Successful requests:                     100       
Failed requests:                         0         
Benchmark duration (s):                  11.85     
Total input tokens:                      117600    
Total generated tokens:                  1000      
Request throughput (req/s):              8.44      
Output token throughput (tok/s):         84.41     
Peak output token throughput (tok/s):    308.00    
Peak concurrent requests:                100.00    
Total token throughput (tok/s):          10011.11  
---------------Time to First Token----------------
Mean TTFT (ms):                          5750.25   
Median TTFT (ms):                        5872.34   
P99 TTFT (ms):                           10549.47  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          536.50    
Median TPOT (ms):                        531.89    
P99 TPOT (ms):                           807.10    
---------------Inter-token Latency----------------
Mean ITL (ms):                           482.85    
Median ITL (ms):                         640.96    
P99 ITL (ms):                            809.80    
==================================================

• three

============ Serving Benchmark Result ============
Successful requests:                     100       
Failed requests:                         0         
Benchmark duration (s):                  11.86     
Total input tokens:                      117600    
Total generated tokens:                  1000      
Request throughput (req/s):              8.44      
Output token throughput (tok/s):         84.35     
Peak output token throughput (tok/s):    308.00    
Peak concurrent requests:                100.00    
Total token throughput (tok/s):          10004.17  
---------------Time to First Token----------------
Mean TTFT (ms):                          5755.36   
Median TTFT (ms):                        5874.58   
P99 TTFT (ms):                           10556.14  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          536.73    
Median TPOT (ms):                        532.49    
P99 TPOT (ms):                           807.43    
---------------Inter-token Latency----------------
Mean ITL (ms):                           483.05    
Median ITL (ms):                         642.68    
P99 ITL (ms):                            810.22    
==================================================

Test Result

---

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.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[lengrongfu]

@mgoin @zyongye Hi, can you help review this pr.

[lengrongfu]

@mgoin Hi, this pr have any issue?

[yewentao256]

Thanks for the work!

Please attach with full reproduce command and full e2e benchmark result logs (running 3 times). Also, please test using lm_eval to make sure acc is correct.

[lengrongfu]

@yewentao256 hi, thanks review, relevant information has been added as requested.

[yewentao256]

Total throughput | 4031.74 tok/s | 3572.61 tok/s | +11.39% -- | -- | -- | --

Why total throughput reduce instead?

[yewentao256]

Please check if this would be safe,

vllm/vllm/model_executor/layers/linear.py

Lines 386 to 394 in 4d30c51

	def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
	if len(loaded_weight.shape) == 0:
	loaded_weight = loaded_weight.reshape(1)
	assert param.size() == loaded_weight.size(), (
	f"Tried to load weights of size {loaded_weight.size()}"
	f"to a parameter of size {param.size()}"
	)
	param.data.copy_(loaded_weight)

[lengrongfu]

In my view, although it is possible to directly overwrite the tensor's memory, vLLM likely does not perform such operations during runtime; I wonder if there are any scenarios I haven't considered.

[lengrongfu]

I can implement this weight splitting more safely within process_weights_after_loading.

[lengrongfu]

Total throughput | 4031.74 tok/s | 3572.61 tok/s | +11.39% -- | -- | -- | --
Why total throughput reduce instead?

There was an error in the compilation of this data; I have provided the latest test results and raw data.

[yewentao256]

Thanks for the work!
Please take a look these AI generated comments

### [P1] Preserve HPC weight storage across weight reloads

`process_weights_after_loading()` reassigns `hpc_w_high`, `hpc_w_low`, and `hpc_split_flag` every time it runs. This changes their storage addresses, while CUDA Graph capture and weight reload flows rely on those addresses remaining stable. In a minimal reproduction, all three pointers changed after reprocessing, so a captured graph may continue reading the old router weights.

Please allocate these tensors only once and use `copy_()` on subsequent reloads. It would also be good to add a weight-reload test with CUDA Graph capture.

[[Code location](https://github.com/vllm-project/vllm/blob/d69c8f8af64e626b0031c6e140c8e2348c21d458/vllm/model_executor/layers/fused_moe/router/gate_linear.py#L195-L205)](https://github.com/vllm-project/vllm/blob/d69c8f8af64e626b0031c6e140c8e2348c21d458/vllm/model_executor/layers/fused_moe/router/gate_linear.py#L195-L205)

### [P1] Guard against split-K workspace overflow

`get_gemm_bf16xfp32_workspace(n)` uses the default `max_tokens=131072`, but the forward path does not enforce any corresponding token limit. For larger `max_num_batched_tokens`, the kernel indexes the workspace using the actual tile/token count, which may result in an out-of-bounds access.

Please size the workspace from the scheduler configuration, or track its capacity and fall back when `x.shape[0]` exceeds it. A boundary test above 131072 tokens would help cover this case.

[[HPC workspace API](https://github.com/Tencent/hpc-ops/blob/main/hpc/gemm.py)](https://github.com/Tencent/hpc-ops/blob/main/hpc/gemm.py) · [[Kernel indexing](https://github.com/Tencent/hpc-ops/blob/main/src/gemm/sm90/gemm_bf16xfp32.cu)](https://github.com/Tencent/hpc-ops/blob/main/src/gemm/sm90/gemm_bf16xfp32.cu)

### [P2] Recompute HPC eligibility in `set_out_dtype()`

`GateLinear` explicitly supports setting `out_dtype` after construction, but HPC eligibility is computed only in `__init__`. If `out_dtype` is initially `None`, `allow_hpc_router_gemm` remains `False` even after calling `set_out_dtype(torch.float32)`. I reproduced this behavior.

Please recompute the HPC eligibility in `set_out_dtype()`, ideally through a shared helper used by both initialization and the setter, and add coverage for the deferred-`out_dtype` path.

[[Code location](https://github.com/vllm-project/vllm/blob/d69c8f8af64e626b0031c6e140c8e2348c21d458/vllm/model_executor/layers/fused_moe/router/gate_linear.py#L150-L193)](https://github.com/vllm-project/vllm/blob/d69c8f8af64e626b0031c6e140c8e2348c21d458/vllm/model_executor/layers/fused_moe/router/gate_linear.py#L150-L193)

[lengrongfu]

@yewentao256 The issues raised in the comments have been addressed.

[yewentao256]

Would this be faster than the default one? We should make it by default if it could work better in any case

[mergify]

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

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