[Core] Use FlashInfer workspace sizing helper by lesj0610 · Pull Request #46883 · vllm-project/vllm

lesj0610 · 2026-06-26T23:43:16Z

Purpose

Use FlashInfer's caller-owned workspace sizing helper when available so vLLM can reserve the workspace size required by FlashInfer attention planning instead of always reserving the default FlashInfer workspace size.

This keeps the existing fallback path intact: if the helper is unavailable or cannot size a particular shape, vLLM still uses FlashInfer's default workspace allocation behavior. The workspace state shared by FlashInfer wrappers now keeps only weak references to wrappers, so profiling-time or temporary wrappers do not stay alive solely because the shared workspace state saw them once.

This PR requires a FlashInfer build that includes flashinfer-ai/flashinfer#3741 for the caller-owned workspace sizing helper. That FlashInfer change has been merged to FlashInfer main, but it is not part of the 0.6.13 release. Until the first FlashInfer release containing it is available and the vLLM pins are bumped to that version, this PR should be tested with a FlashInfer main checkout/build.

AI assistance: Codex was used during implementation and PR preparation; the submitter reviewed the changes.

Test Plan

Run syntax, lint, and format checks for the modified FlashInfer backend and CUDA graph profiling tests.
Run focused CUDA graph memory profiling tests with a FlashInfer checkout that provides the workspace sizing helper.
Let the repository commit hooks run for the cleanup commits, including ruff, format, mypy, SPDX, and repository policy checks.
After this branch is integrated with the serving branch, run Qwen3.6 long-context nvfp4 serving smoke to verify the workspace reserve path before CUDA graph lock.

Test Result

Syntax, lint, format, and whitespace checks:

python -m py_compile \
  vllm/v1/attention/backends/flashinfer.py \
  tests/v1/worker/test_cudagraph_memory_profiling.py

python -m ruff check \
  vllm/v1/attention/backends/flashinfer.py \
  tests/v1/worker/test_cudagraph_memory_profiling.py

python -m ruff format --check \
  vllm/v1/attention/backends/flashinfer.py \
  tests/v1/worker/test_cudagraph_memory_profiling.py

git diff --check

Result: passed.

Dependency version checks:

.venv/bin/python -m json.tool docker/versions.json
git grep -n "0\.6\.12" -- requirements docker tools pyproject.toml
git grep -n "flashinfer-python==0\.6\.13\|flashinfer-cubin==0\.6\.13\|FLASHINFER_VERSION=0\.6\.13\|v0\.6\.13" -- requirements docker tools pyproject.toml

Result: JSON validation passed and no FlashInfer 0.6.12 references remain in the checked requirement and Docker paths. The checked paths still pin FlashInfer 0.6.13; those pins need to move to the first FlashInfer release that includes flashinfer-ai/flashinfer#3741 once it is published. Before that release, testing should use a FlashInfer main checkout/build.

Focused CUDA graph profiling test:

PYTHONPATH=<flashinfer checkout> FLASHINFER_DISABLE_VERSION_CHECK=1 \
  python -m pytest tests/v1/worker/test_cudagraph_memory_profiling.py -q

Result before the latest reserve-path test additions: 5 passed.

After integrating this branch into the serving integration branch, the full CUDA graph memory profiling test file was also run in the built vLLM development environment:

python -m pytest tests/v1/worker/test_cudagraph_memory_profiling.py -q

Result: 7 passed.

Qwen3.6 long-context nvfp4 serving smoke after integration:

Model: Qwen3.6-27B AWQ, tensor parallel size 2, --kv-cache-dtype nvfp4, --max-model-len 32768, --max-num-seqs 64.
FlashInfer workspace was reserved before CUDA graph capture: 0.00 MB -> 61.74 MB, then locked at 61.740234375 MB.
KV capacity: 834,491 tokens.
CUDA graph memory comparison: 0.83 GiB actual pool, 0.85 GiB estimated pool, 0.08 GiB persistent estimate.
Random serving smoke: 16 successful requests, 0 failed requests.
MRCR smoke: 30 samples, match ratio 0.9202, prefix hit rate 1.0000, no workspace grow failure or engine fatal error in the server log.

Commit hooks passed for the latest cleanup commit, including ruff, format, mypy, SPDX, config validation, and related repository checks.

Not yet verified in this branch:

Full Gemma4 serving capacity smoke after integration with the serving branch.
Multilingual generation-quality smoke after integration with the serving branch.

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. Not needed; this is an internal FlashInfer workspace accounting change.