[abmfy]

Purpose

With --linear-backend flashinfer_cutedsl, every unquantized BF16 linear goes through the opaque vllm.cuda_flashinfer_bf16_gemm op. FlashInfer runs its CuTeDSL kernel only for M ≤ 32; since #60434, larger M goes to FlashInfer's autotuned cuBLASLt. On steps above the CUDA graph capture size, each of those calls still costs 145–167 µs of host time: the dispatcher, the eligibility checks, FlashInfer's Python path and Inductor's output asserts. F.linear costs about 20 µs. Prefill steps become host-bound, and TTFT rises by 34–48% (table below).

This PR:

• Splits the compile ranges at 32 for this backend, and adds a post-grad pass that replaces the op with aten linear in the ranges above 32. Inductor then emits the same extern mm as the default backend.
• Sends M > 32 to F.linear outside torch.compile, i.e. with enforce_eager and inside opaque ops such as moe_forward_shared.
• Adds a 32-token FlashInfer autotune pass. The max_num_batched_tokens pass no longer reaches the op, so the CuTeDSL buckets would otherwise never be tuned.
• Gates all of this on KernelConfig.use_flashinfer_cutedsl_bf16_gemm.

At these shapes, FlashInfer's tuned cuBLASLt and aten mm take the same GPU time. Measured with CUPTI inside CUDA graphs, on 14 Qwen3.5 linear shapes at M = 64–16384: the per-forward ratio is 0.96–1.01, and at M = 1031 both pick the same nvjet kernel.

Test Plan

GB200 (SM100), CUDA 13, torch 2.13, FlashInfer 0.7.1 (#60948). The serving numbers compare main 258b97ff6b with this change on top of it; the branch has since been rebased onto ec317a4618 without conflicts, and the unit tests below were re-run there.

Each model runs on one node: a cold round, then 3 interleaved warm rounds. All arms are seeded with one FlashInfer autotune file.

pytest tests/compile/passes/test_flashinfer_bf16_gemm_fallback.py tests/model_executor/test_flashinfer_autotune_warmup.py tests/kernels/test_flashinfer_bf16_gemm.py
vllm serve <model> --max-model-len 32768 --reasoning-parser qwen3 --enable-prefix-caching [--linear-backend flashinfer_cutedsl]
vllm bench serve --backend openai-chat --dataset-name random --temperature 0 --ignore-eos \
  --random-input-len 1000 --random-output-len 1000 --max-concurrency {1,32}   # w1, w2
  # w5: --random-input-len 16384 --random-output-len 128 --max-concurrency 8
python tests/evals/gsm8k/gsm8k_eval.py --num-questions 1319 --num-shots 5

Test Result

The new pass test and the updated warmup test fail on main and pass here.

TP1, mean over 3 rounds. "flag" means --linear-backend flashinfer_cutedsl.

	35B base	35B main + flag	35B PR + flag	4B base	4B main + flag	4B PR + flag
c=1 out tok/s	312.7	369.2	365.5	340.4	389.7	389.7
c=32 mean TTFT ms	376.2	529.5 (+40.8%)	374.9 (−0.3%)	270.5	361.7 (+33.7%)	272.3 (+0.7%)
16k/128 c=8 mean TTFT ms	836.8	1234.0 (+47.5%)	844.3 (+0.9%)	816.4	820.7	814.5
16k/128 c=8 out tok/s	600.8	477.2	611.5	660.0	675.3	676.7
GSM8K	0.861	0.867	0.862	0.763	0.776	0.772
warm restart, s	401	468	387	344	423	351

35B is Qwen3.5-35B-A3B and 4B is Qwen3.5-4B.

• c=1: on 35B the PR is −1.0% vs main + flag, 95% CI [−5.6, +3.5], driven by one slow round. On 4B the two are identical.
• Restarts: main + flag re-tunes its M > 32 cuBLASLt entries on every start (52–67 s), because FlashInfer's cache key is the exact M. The PR hits the autotune cache on every warm start.
• Greedy outputs: on 35B they change across server restarts in every arm, including base, so this is not caused by the PR. On 4B they are identical across rounds, and the PR matches main + flag on all 16 prompts.

Not a duplicate: I found no open PR that changes the dispatch of cuda_flashinfer_bf16_gemm or adds per-range linear passes. This complements #60434. If #59021 lands, use_flashinfer_cutedsl_bf16_gemm should read the effective unquantized backend. #59244 edits the same autotune token-count function.

AI assistance (Claude Code) was used for the code, the tests, the benchmarks and this description.

🤖 Generated with Claude Code

[mergify]

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

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