[Kernel] Warm up hybrid GDN/Mamba/MRoPE kernels

[lesj0610]

Purpose

Hybrid Qwen-style GDN/Mamba/MRoPE models can trigger Triton compilation during the first inference request when startup warmup does not cover the exact runtime compile keys. This revision migrates the remaining warmups to the shared kernel-owned warmup contract and removes coverage that is already provided by current upstream main.

Changes

• Merge the latest upstream main and remove duplicate warmups for zero-KV blocks, slot mapping, causal convolution, fused post-convolution, fused sigmoid update, and chunked GDN prefill.
• Add kernel-owned VllmJitKernel implementations for packed recurrent GDN decode and MRoPE.
• Route runtime launches and warmup compilation through the same owners and exact CompileKey definitions.
• Use compile-only warmup calls under enable_jit_warmup; no dummy model execution or startup synchronization is added.
• Cover the integer specialization buckets used by MRoPE and the static meta-parameters used by packed GDN decode.

Test Plan

.venv/bin/python -m pytest \
  tests/model_executor/test_hybrid_gdn_mamba_mrope_warmup.py -q

CUDA_VISIBLE_DEVICES=1 .venv/bin/python -m pytest \
  tests/kernels/core/test_mrope.py -q

pre-commit run --files \
  vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py \
  vllm/model_executor/layers/rotary_embedding/mrope.py \
  vllm/model_executor/warmup/hybrid_gdn_mamba_mrope_warmup.py \
  vllm/model_executor/warmup/kernel_warmup.py \
  vllm/third_party/flash_linear_attention/ops/fused_recurrent.py \
  tests/model_executor/test_hybrid_gdn_mamba_mrope_warmup.py

GPU compile-only and direct numerical smoke tests were run on GPU 1 (RTX 3090).

Test Result

• Shared-contract unit tests: 3 passed.
• MRoPE native-versus-CUDA tests: 12 passed.
• Compile-only smoke compiled both MRoPE keys and the packed recurrent GDN key.
• Direct packed recurrent GDN runtime comparison had maximum absolute error 0.007812, within the existing test tolerance.
• All applicable pre-commit hooks, including mypy, passed.

The existing packed recurrent numerical pytest still fails for fp16/bf16 on both this branch and an unmodified latest-upstream worktree; fp32 passes. This is therefore not introduced by this revision.

AI assistance: Codex.

[gemini-code-assist]

Code Review

This pull request introduces kernel warmup mechanisms for hybrid GDN/Mamba and MRoPE models to prevent JIT compilation overhead and potential OOM issues during the first real inference requests. It adds dedicated warmup routines for prefill and decode paths, integrates these into the main kernel warmup pipeline, ensures robust cleanup of KV connector states on failure, and adds comprehensive unit tests. No review comments were provided, so there is no feedback to address.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f92cd20cbc

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[mergify]

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

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

[mergify]

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

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

[mergify]

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

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

[LopezCastroRoberto]

Hey @lesj0610

Similar thing to what I mentioned in another of your PRs. We're currently migrating all kernel warmups to a shared warmup contract. See #47451, RFC: #47456.

Would you mind migrating these warmups to conform to that shared contract? It would help keep the warmup infrastructure consistent and make future maintenance easier.

However, I think a sync with upstream is needed. Some of the kernels warmed up on this PR are already covered upstream, while some of them are not. Thanks!

Adding your PR to the list: #49349

[mergify]

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

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

[mergify]

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

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

[mergify]

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

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