[mmangkad]

Purpose

Fix TRTLLM_RAGGED MLA prefill failing after CUDA graph capture.

After the MLA prefill backend refactor in #32623, the FlashInfer workspace used by TRTLLM_RAGGED MLA prefill is allocated lazily. The first real request after lock_workspace() could hit:

Workspace is locked but allocation from 'trtllm_ragged.py:66:_get_workspace_buffer' requires 394.00 MB, current size is 0.00 MB. Workspace growth is not allowed after locking.

Eagerly allocate the FlashInfer workspace buffers during MLA prefill backend initialization so they are reserved before lock_workspace().

cc @MatthewBonanni @LucasWilkinson @mgoin

Test Plan

Run the server with TRTLLM_RAGGED MLA prefill on GB300 GPUs:

vllm serve moonshotai/Kimi-K2.6 \
  --trust-remote-code \
  --tensor-parallel-size 4 \
  --attention-config.mla_prefill_backend TRTLLM_RAGGED \
  --tool-call-parser kimi_k2 \
  --enable-auto-tool-choice \
  --reasoning-parser kimi_k2 \
  --speculative-config '{"model": "lightseekorg/kimi-k2.6-eagle3-mla", "method": "eagle3", "num_speculative_tokens": 3}'

Test Result

Before this PR, the command above did not start successfully and failed with the locked workspace assertion.

After this PR, the same command starts successfully and works correctly with TRTLLM_RAGGED MLA prefill.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[gemini-code-assist]

Code Review

This pull request updates the _warmup_and_capture logic in gpu_model_runner.py to ensure force_attention is enabled for all captured CUDA graph modes, rather than only for full graphs. This adjustment ensures that attention metadata is built during warmup, preventing forbidden lazy workspace allocations after lock_workspace() is called. I have no feedback to provide as no review comments were submitted.

[mergify]

Hi @mmangkad, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mmangkad]

@MatthewBonanni @LucasWilkinson could you help take a look at this?

[MatthewBonanni]

Thanks for finding this bug and contributing this fix! I'd prefer to just eagerly allocate the buffers though, rather than force_attention for PIECEWISE warmups.

#42539

[mmangkad]

@MatthewBonanni updated

[MatthewBonanni]

LGTM, thanks!