[Bugfix] limit cudagraph capture sizes by num_blocks for GDN models

[ZJY0516]

Purpose

FIX #34094
Fixes a corner case bug where vLLM crashes with an AssertionError during CUDA graph capture for GDN (Gated Delta Net) models when the configured cudagraph capture size exceeds the available KV cache blocks (num_blocks).

Problem

For GDN models using CUDA graphs, vLLM determines the cudagraph capture sizes at config initialization time based on max_num_seqs. However, the actual number of KV cache blocks (num_blocks) is determined later during memory profiling and can be smaller than the cudagraph capture size due to memory constraints.

When this happens, during causal_conv1d_update, the following assertion fails because num_cache_lines (which equals num_blocks) is smaller than the batch size:

assert num_cache_lines >= batch  # AssertionError!

The GDN attention backend uses num_blocks as cache lines for conv states. When CUDA graph capture creates batches larger than num_blocks, the conv state cache doesn't have enough slots, triggering the assertion.

Solution

Add _maybe_limit_cudagraph_sizes_by_num_blocks() in GPUModelRunner.initialize_kv_cache() to:

Check if the model uses GDN attention (by detecting GDN_ATTN backend)
Check if CUDAGraphMode has FULL mode enabled (only affects FULL/FULL_AND_PIECEWISE modes)
If max_cudagraph_capture_size > num_blocks, filter cudagraph_capture_sizes to only include sizes ≤ num_blocks and update max_cudagraph_capture_size accordingly
This ensures CUDA graph capture never attempts batch sizes larger than available cache lines for GDN models.

Test

vllm serve Qwen/Qwen3.5-35B-A3B --port 8001 --gpu-memory-utilization 0.55

log

(EngineCore_DP0 pid=808156) WARNING 02-27 16:51:15 [gpu_model_runner.py:6032] Limiting max_cudagraph_capture_size from 512 to 336 due to num_blocks=348 constraint for GDN model

lm_eval \
    --model local-completions \
    --model_args "model=Qwen/Qwen3.5-35B-A3B,base_url=http://0.0.0.0:8001/v1/completions,tokenized_requests=False,tokenizer_backend=None,num_concurrent=256,timeout=5000,max_length=4096" \
    --tasks gsm8k \
    --num_fewshot 5

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.8658	±	0.0094
		strict-match	5	exact_match	↑	0.8484	±	0.0099

cc @tdoublep @ywang96

---

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.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request addresses an assertion error by improving the validation logic in causal_conv1d_update. The new assertion correctly checks that all conv_state_indices are within the valid range of num_cache_lines. However, I've identified a critical edge case in the new code. When the batch size is zero, conv_state_indices.max() will be called on an empty tensor, causing a RuntimeError. I've provided a suggestion to fix this.

[tdoublep]

I don't really understand this change. If the assert is duplicated - why doesn't the first one fail?

This check seems like a correct thing to have imo.

[ZJY0516]

I don't really understand this change. If the assert is duplicated - why doesn't the first one fail?

This check seems like a correct thing to have imo.

Because conv_state_indices is not None in this case

[tdoublep]

OK, but then we read the batch size from the conv_state_indices, and we want to verify that the number of possible blocks (e.g., the outer dimension of conv_state) is greater than or equal to the batch size. Isn't this a reasonable check? Is this something to do with the interaction of DP?

[ZJY0516]

OK, but then we read the batch size from the conv_state_indices, and we want to verify that the number of possible blocks (e.g., the outer dimension of conv_state) is greater than or equal to the batch size. Isn't this a reasonable check? Is this something to do with the interaction of DP?

Yes, you are right. I'll try to fix it in another way

[mergify]

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

uv pip install pre-commit
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 or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

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

[tdoublep]

PTAL at this other PR that seems to implement similar solution:
#34571

[tdoublep]

Why is this a specific change to GDN? I think we need to do this for all hybrid models actually.

[tdoublep]

What happens if this is empty?

[tdoublep]

Does this need to be in GPU model runner?