[amd-dlimpus]

Replaces #64264, which GitHub closed when its base branch minimax-m31-flash was deleted after #64111 merged into main.

Overview

With an NVFP4 KV cache on ROCm (gfx950), each FlyDSL launcher compiles its kernel variant on first launch, and the variant depends on the batch and the context length. Nothing compiles them ahead of time: minimax_m3_msa_warmup prebuilds kernels on Blackwell only, and vLLM's JIT monitor only sees Triton. So every new variant compiles inside a request. This PR adds a ROCm warmup that compiles all of them at startup.

Claims

• On a fresh MI355X TP4 server (NVFP4 KV, Quark MXFP4 experts, 8K in / 512 out), the first C16 run no longer pays the FlyDSL compile. Mean TTFT goes from 3143 ms to 1691 ms, and p99 from 6471 ms to 3369 ms. That matches FP8 KV's first run (1694 ms). The rest of the cold cost is Triton JIT that both KV formats share, which this PR doesn't touch.
• Warm throughput is unchanged; only first-use compiles move to startup. This adds 11.6 s per rank to startup; the ranks warm up in parallel.
• Only the ROCm NVFP4 KV path is affected. Other ROCm models return at a sys.modules check, and CUDA is untouched.

Validation

Root cause. Timed on one MI355X: the first call of minimax_m3_nvfp4_index_score_prefill takes 3297 ms for two 8K sequences (the large kernel) and 796 ms for one (kernel_shared). Warm calls take under 0.1 ms. A single 8K prefill uses the shared kernel, so C1 and the usual warmup never reach the large one. The first step with two 8K prompts, at C16, compiles it mid-request. Before the fix, C16 mean TTFT was 2849 ms with NVFP4 KV against 1475 ms with FP8 KV, and p99 was +3 s; C1 and C64 matched.

Test. tests/kernels/attention/test_minimax_m3_nvfp4_warmup.py runs the warmup at given scheduler limits, then makes 40 random in-range launches of every launcher. It asserts that no launcher's kernel cache grows. It passes on MI355X with 1 and 4 index heads. Mutation check: with the large-prefill sweep removed, the test fails on a random draw (39 requests, 5867 tokens).

CI. The test sits with #63232's other FlyDSL NVFP4 kernel tests. The sharded kernels/attention steps collect it but skip it, because no Buildkite step runs on gfx950 (the MI300 shards fail the gfx950 check). It was validated on MI355X outside CI.

End to end. MI355X, TP4, Quark MXFP4 experts. Each server boots fresh and runs back-to-back C16 passes of 8K in / 512 out, over 32 text prompts.

Server	First pass mean / p99 TTFT	First pass tok/s	Second pass mean TTFT	Second pass tok/s
NVFP4 KV, before	3143 / 6471 ms	826	340 ms	1482
NVFP4 KV, this PR	1691 / 3369 ms	975	505 ms	1464
FP8 KV (reference)	1694 / 3409 ms	980	435 ms	1402

With this PR, the only compiles the JIT monitor reports during inference are the four Triton kernels the FP8 KV server also compiles (_mxfp8_linear_kernel, _reduce_terms_kernel, _gemm_a16_w16_kernel, _prepare_pos_seq_lens_kernel). Each server ran once, so second-pass TTFTs vary by about ±100 ms.

Details

Variants are keyed per launcher as follows:

• fused norm / RoPE / insert: tokens ≤ 512 or above.
• sparse prefill: one or two combine rounds.
• sparse decode: split count, which depends on rows.
• index prefill: the shared or the large kernel.
• index decode: rows, and whether the score is buffered, which depends on requests × context.
• top-k: strip count, which depends on context, plus the decode kernel.

The warmup drives each public launcher over exactly those inputs, bounded by max_num_seqs, max_num_batched_tokens, max_model_len and the speculative query lengths. It uses the layer's caches and throwaway outputs: inserts get slot -1, and the other kernels only read. FULL decode graphs already fix their variants at capture. The sweeps exist for prefill and mixed steps, which run these launchers eagerly.

The decode-CP push kernel is not swept. It has a single variant and is compiled by the FULL decode-graph capture.

Persisting FlyDSL's cache (~/.flydsl/cache in the container) across restarts would also remove the startup cost after the first boot. That is a deployment choice and not part of this PR.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[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.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀