[UX] Speedup DeepGEMM warmup with heuristics

[mgoin]

Purpose

By looking at get_best_config in DeepGEMM's https://github.com/deepseek-ai/DeepGEMM/blob/79f48ee15a82dd5fad5cd9beaa393c1f755e6b55/csrc/jit_kernels/heuristics/common.hpp we realized that we can greatly reduce the number of batch sizes we have to run for each weight shape when precompiling DeepGEMM. This greatly improves startup time when hitting warm DeepGEMM cache because we need to run many fewer gemms.

Test Plan

Test that all the kernels necessary are generated before and after this PR

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.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request introduces a heuristic-based approach to speed up the DeepGEMM warmup process. By selectively choosing M values for GEMM operations instead of iterating through all possibilities, it significantly reduces the number of warmup iterations. The changes look good, but there is a significant issue with hardcoding device 0 when querying for device properties, which could lead to incorrect warmup behavior in multi-GPU environments. I've added comments with suggestions to address this.

[varun-sundar-rabindranath]

Nice changes @mgoin!

I verified with command
vllm bench throughput --model deepseek-ai/DeepSeek-V3 --hf-overrides.num_hidden_layers=4 --load-format=dummy --enforce-eager --input-len=10000 --output-len=500 --num-prompts=1
that the kernels generated are the same between main and this PR.

Verified the same for Qwen/Qwen3-30B-A3B-FP8

[varun-sundar-rabindranath]

Like we talked over chat - I think the default should be processing all the token sizes so we can be sure that in deployment scenarios we aren't jitting in the hotpath. We can introduce an env var VLLM_RELAX_DEEP_GEMM_WARMUP and hide _generate_optimal_warmup_m_values behind that (like something people can opt-in).
Thanks 🙌

[varun-sundar-rabindranath]

LGTM! Thanks @mgoin

[varun-sundar-rabindranath]

@tlrmchlsmth @yewentao256 can you please take a look when you find some time. Thanks 🙌

[yewentao256]

Thanks for the work!

I am thinking if we add too many env variables and make it complicated to use.
Would full warm up at the first time and we automatically skip second warm up if cache is there work?

[yewentao256]

from vllm.utils import cdiv

[varun-sundar-rabindranath]

Would full warm up at the first time and we automatically skip second warm up if cache is there work?

@yewentao256 We considered this and decided against it for a first pass as it is hard to detect if the cache has a JIT for a particular GEMM shape.

1. About the detection logic itself, it is insufficient to use just the model name as key, we would also need to consider parallelization configs to generate a reasonable key. This is can get unruly. [we could probably reuse the hash that we generate for torch compile - I thought of this as I was typing and could be useful]
2. Updates in deep_gemm: If deep_gemm updates and we skip warmup due to some local detection logic, we will run with old kernels. [Maybe we can add the deep_gemm commit hash to torch compile hash]

I am thinking of VLLM_RELAX_DEEP_GEMM_WARMUP as an env var for power users. For normal users, I think the default recommendation should be to use neither VLLM_RELAX_DEEP_GEMM_WARMUP nor VLLM_SKIP_DEEP_GEMM_WARMUP as it guarantees optimal performance.

[tlrmchlsmth]

is it too dangerous to default to relaxed?