[Misc] DeepGEMM : Avoid JIT generation in the hot-path

[varun-sundar-rabindranath]

Purpose

DeepGemm JITs the its GEMM kernels. This JIT generation depends on the Gemm input tensor shapes. We want to avoid generating JITs during actual model inference. To this effect, this PR introduces a "Kernel Warmup" phase during which all the DeepGEMM kernels are invoked with possible input shapes for the model so all JIT generation is complete.

Test Plan

VLLM_ALL2ALL_BACKEND="deepep_high_throughput" VLLM_USE_DEEP_GEMM=1 canhazgpu run -g2 -- vllm serve Qwen/Qwen3-30B-A3B-FP8  --trust-remote-code --enable-expert-parallel --data-parallel-size 2 --port 9010 --no-enable-prefix-caching

Notes:

• Cleared the cache before executing the vllm serve command.
• The fp8_gemm_nt warmup took 83s
• The m_grouped_fp8_gemm_nt_contiguous warmup took 48s

Test Result

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  | 0.84|±  |0.0368|
|     |       |strict-match    |     5|exact_match|↑  | 0.93|±  |0.0256|

Performance

IMPORTANT: Note that numbers are better with this PR. This simply means that main in steady-state (i.e. when run long enough to have all the kernels JIT'ed) will also produce these numbers.

VLLM_ALL2ALL_BACKEND="deepep_high_throughput" VLLM_USE_DEEP_GEMM=1 canhazgpu run -g2 -- vllm serve Qwen/Qwen3-30B-A3B-FP8  --trust-remote-code --enable-expert-parallel --data-parallel-size 2 --port 9010 --no-enable-prefix-caching

python3 benchmarks/benchmark_serving.py --model  Qwen/Qwen3-30B-A3B-FP8 --dataset-name sharegpt --dataset-path ./ShareGPT_V3_unfiltered_cleaned_split.json --num-prompts 1000 --port 9010

main - Does m_grouped_fp8_gemm_nt_contiguous warmup only
============ Serving Benchmark Result ============
Successful requests:                 	1000 	 
Benchmark duration (s):              	100.87    
Total input tokens:                  	217393    
Total generated tokens:              	201847    
Request throughput (req/s):          	9.91 	 
Output token throughput (tok/s):     	2001.14   
Total Token throughput (tok/s):      	4156.40   
---------------Time to First Token----------------
Mean TTFT (ms):                      	2771.34   
Median TTFT (ms):                    	2656.45   
P99 TTFT (ms):                       	6170.29   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                      	402.21    
Median TPOT (ms):                    	213.52    
P99 TPOT (ms):                       	1851.69   
---------------Inter-token Latency----------------
Mean ITL (ms):                       	188.65    
Median ITL (ms):                     	85.00	 
P99 ITL (ms):                        	5031.81   
==================================================

PR:
============ Serving Benchmark Result ============
Successful requests:                 	1000 	 
Benchmark duration (s):              	71.35	 
Total input tokens:                  	217393    
Total generated tokens:              	201847    
Request throughput (req/s):          	14.02	 
Output token throughput (tok/s):     	2828.97   
Total Token throughput (tok/s):      	5875.82   
---------------Time to First Token----------------
Mean TTFT (ms):                      	2532.40   
Median TTFT (ms):                    	2579.44   
P99 TTFT (ms):                       	3738.82   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                      	106.27    
Median TPOT (ms):                    	93.48	 
P99 TPOT (ms):                       	212.56    
---------------Inter-token Latency----------------
Mean ITL (ms):                       	89.68	 
Median ITL (ms):                     	87.83	 
P99 ITL (ms):                        	176.79    
==================================================

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[gemini-code-assist]

Code Review

This pull request refactors the DeepGEMM kernel warmup logic to avoid JIT compilation in the hot path of model inference. It introduces a dedicated warmup phase during model initialization. The changes are well-structured and correctly move the warmup logic. My main feedback is to improve the performance of one of the warmup loops, which could be excessively slow for large max_tokens values.

[varun-sundar-rabindranath]

Trigger for the DeepGEMM warmup kernels.
I have added a new folder vllm/warmup where we can add kernel warmup code going forward.

[mgoin]

I was initially thinking this would be in profile_run or capture_model in the model_runner, but I think this is a better place

[varun-sundar-rabindranath]

yeah. Didn't want to affect profile_run with arbitrary code that might mess up the peak memory usage.

[mgoin]

I was initially thinking this would be in profile_run or capture_model in the model_runner, but I think this is a better place

[bnellnm]

Could you use _valid_deep_gemm_shape here?

[varun-sundar-rabindranath]

no. In addition to checking the weight dimensions _valid_deep_gemm_shape employs runtime heuristics on the M dimension which we don't have here.

[bnellnm]

deep gemm can also be called by the triton fused_experts under the right conditions. not sure how easy that is to check though.

[varun-sundar-rabindranath]

deep gemm can also be called by the triton fused_experts under the right conditions
Hey @bnellnm I not sure which invocation you are referring to. Can you point me at it please.

If you are referring to TritonOrDeepGemmExperts ? I have a condition to handle below this.

[bnellnm]

fused_experts in fused_moe.py at line ~1290 it can call deep_gemm_moe_fp8 if the flags+size are set properly. It might not be easy to check for this though.

[varun-sundar-rabindranath]

Oh I considered that. we return True here to account for that.
if module.quant_method.fused_experts is not a modular kernel, then we assume that we might be invoking deep_gemm kernels and conservatively do the warmup.

About checking for the exact set of conditions, it is hard to keep them in sync and I think it is instead better to default to doing the warmup.

• The actual jitting during the warmups is done only once per model. It should complete in a reasonable time. But if it is undesirable, users can always opt-out with the VLLM_SKIP_DEEP_GEMM_WARMUP flag.
• On subsequent runs, this loop should be fast as it'd just be only executing a series kernels.

what do you think ?

[bnellnm]

Ok, that makes sense.

[varun-sundar-rabindranath]

Hi @zou3519 . The changes to this file is due to the lint-and-deploy CI job failing with,

(EngineCore_0 pid=34)   File "/opt/venv/lib/python3.12/site-packages/vllm/model_executor/layers/fused_moe/fused_moe.py", line 1060, in <module>
(EngineCore_0 pid=34)     direct_register_custom_op(
(EngineCore_0 pid=34)   File "/opt/venv/lib/python3.12/site-packages/vllm/utils/__init__.py", line 2491, in direct_register_custom_op
(EngineCore_0 pid=34)     schema_str = torch.library.infer_schema(op_func,
(EngineCore_0 pid=34)                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_0 pid=34)   File "/opt/venv/lib/python3.12/site-packages/torch/_library/infer_schema.py", line 106, in infer_schema
(EngineCore_0 pid=34)     error_fn(
(EngineCore_0 pid=34)   File "/opt/venv/lib/python3.12/site-packages/torch/_library/infer_schema.py", line 58, in error_fn
(EngineCore_0 pid=34)     raise ValueError(
(EngineCore_0 pid=34) ValueError: infer_schema(func): Parameter block_shape has unsupported type typing.Optional[list[int]]. The valid types are: dict_keys([<class 'torch.Tensor'>, typing.Optional[torch.Tensor], typing.Sequence[torch.Tensor], typing.List[torch.Tensor], typing.Sequence[typing.Optional[torch.Tensor]], typing.List[typing.Optional[torch.Tensor]], <class 'int'>, typing.Optional[int], typing.Sequence[int], typing.List[int], typing.Optional[typing.Sequence[int]], typing.Optional[typing.List[int]], <class 'float'>, typing.Optional[float], typing.Sequence[float], typing.List[float], typing.Optional[typing.Sequence[float]], typing.Optional[typing.List[float]], <class 'bool'>, typing.Optional[bool], typing.Sequence[bool], typing.List[bool], typing.Optional[typing.Sequence[bool]], typing.Optional[typing.List[bool]], <class 'str'>, typing.Optional[str], typing.Union[int, float, bool], typing.Union[int, float, bool, NoneType], typing.Sequence[typing.Union[int, float, bool]], typing.List[typing.Union[int, float, bool]], <class 'torch.dtype'>, typing.Optional[torch.dtype], <class 'torch.device'>, typing.Optional[torch.device]]). Got func with signature (hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, activation: str = 'silu', is_act_and_mul: bool = True, apply_router_weight_on_input: bool = False, use_fp8_w8a8: bool = False, use_int8_w8a8: bool = False, use_int8_w8a16: bool = False, use_int4_w4a16: bool = False, use_mxfp4_w4a4: bool = False, per_channel_quant: bool = False, global_num_experts: int = -1, expert_map: Optional[torch.Tensor] = None, w1_scale: Optional[torch.Tensor] = None, w2_scale: Optional[torch.Tensor] = None, w1_zp: Optional[torch.Tensor] = None, w2_zp: Optional[torch.Tensor] = None, a1_scale: Optional[torch.Tensor] = None, a2_scale: Optional[torch.Tensor] = None, block_shape: Optional[list[int]] = None) -> None)

example failing instance https://github.com/vllm-project/vllm/actions/runs/16761346188/job/47457097769?pr=22215

I did some debugging and it looks like the failure is triggered by the inspection of the fused_experts object in file deep_gemm_warmup.py and function _fused_moe_grouped_gemm_may_use_deep_gemm (Both added in this PR). I am not sure why the error wasn't triggered before. However, is this the right fix ? Appreciate any comments / thoughts to make this better. Thanks 🙌

[zou3519]

The fix is good. The behavior change is weird, infer_schema should accept both list and List. Is the lint job using the right version of PyTorch? (needs >= 2.6)

[zou3519]

Okay, the problem is that infer_schema doesn't support Optional[list[Tensor]], but it does support Optional[List[Tensor]]. infer_schema supports list[] in some other situations.

That being said, I don't know why the logic in this PR changed. It sounds like we were not calling infer_schema on inplace_fused_experts before, but now we are. Do you have a sense of why?

[varun-sundar-rabindranath]

It is not just inplace_fused_experts , but also outplace_fused_experts and flashinfer_fused_moe_blockscale_fp8. I had to use List in all of those prototype.

Like I mentioned before, it looks like inspecting module.quant_method.fused_experts triggers torch.library.infer_schema . Also, note that I don't run into this issue locally (using Nvidia GPUs), it seems to happen only the lint-and-deploy job that uses the vllm-cpu docker image.