[Fangzhou-Ai]

Overview

Warm up AITER's FlyDSL GEMM for 32x32-block MXFP8 linears (DeepSeek V4.1 on gfx950) before serving, so a cold server does not JIT-compile FlyDSL kernels inside its first long prefill.

Stacked on #60169, which adds the FlyDSL path. The first commit here is #60169's commit rebased onto current main, with conflicts resolved: it keeps main's pre-quantized QuantizedActivation input on the FlyDSL branch. Please review only the last commit. I'll rebase onto #60169 once it lands, or the author can fold this commit into #60169 instead.

Claims

• After the warmup, serving any token count up to max_num_batched_tokens compiles no FlyDSL kernel. Without it, those compiles happen on the serving path.

Validation

tests/kernels/quantization/test_rocm_mxfp8_linear.py::test_rocm_flydsl_mxfp8_warmup_leaves_nothing_to_compile[warmup|no_warmup] on MI355X:

• Setup: each case runs in a subprocess with an empty FLYDSL_RUNTIME_CACHE_DIR.
• What it checks: it runs a V4.1 TP1 shape over every token count 1..512 and counts compiled kernels.
• Results:
  • with the warmup, the sweep adds 0 compiled kernels;
  • without it, the sweep compiles them;
  • the no_warmup case guards that the test actually exercises the cold path.

pytest tests/kernels/quantization/test_rocm_mxfp8_linear.py -k "warmup or flydsl"
45 passed

The rest of the file (#60169's FlyDSL accuracy tests) also passes on the rebased base commit. The test file is already in the AMD kernels CI list.

Details

• Why it's needed: AITER selects a FlyDSL kernel per padded token count. With AITER's tuned table, each V4.1 MXFP8 shape maps 1..8192 tokens onto 22–41 distinct kernels, about 200 across the six shapes, each JIT-compiled on first use.
  • CUDA-graph capture only compiles the captured sizes (≤ 1024 tokens).
  • The remaining prefill sizes compile during serving.
  • In DeepSeek-V4.1 DPATP2 serving on MI355X, a cold server compiled several inside its first long prefill and hit the execute-model RPC timeout.
• Where it runs: rocm_flydsl_mxfp8_warmup sits in kernel_warmup next to the other kernel warmups (ROCm, enable_jit_warmup). For each FlyDSL linear shape, it runs the GEMM once per distinct kernel that AITER's public get_mxscale_bpreshuffle_config selects for 1..max tokens. Untuned shapes are sampled at powers of two.
• Why not the JIT warmup registry: that registry traces vLLM-owned kernels, while this GEMM is AITER's external op. The same pattern as deep_gemm_warmup fits better.
• Cost: the warmup moves those compiles to startup. With a warm FlyDSL cache it is just one GEMM per distinct kernel.

Duplicate check: no open PR warms the FlyDSL MXFP8 GEMM (searched "flydsl warmup", "warm up mxfp8", "bpreshuffle warmup"). #64264 prebuilds MiniMax-M3 NVFP4 FlyDSL kernels, which is a different model and op. The same patch was offered on #60169 in this comment.

---

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.

AI assistance (Claude) was used for this change; I reviewed it and ran the tests above.

🤖 Generated with Claude Code

[Fangzhou-Ai]

I would recommend to maybe add these JIT cache to the aiter wheel not on vLLM side if possible

[Fangzhou-Ai]

This should be gated as ROCm only

[Fangzhou-Ai]

Is this a universal limit for all FlyDSL GEMM used in vLLM?