[Bugfix] Disable TRTLLM attention when KV transfer is enabled

[ZhanqiuHu]

Summary

• Problem: On Blackwell GPUs, vLLM crashes with AssertionError: is_strictly_contiguous(kv_cache_permute) when using NixlConnector for P/D disaggregation.
• Cause: FlashInfer auto-enables TRTLLM attention on Blackwell, which requires contiguous KV cache tensors. NixlConnector creates non-contiguous KV cache views.
• Fix: Auto-disable TRTLLM attention when kv_transfer_config is set.

Purpose

On Blackwell GPUs, FlashInfer auto-detects and uses TRTLLM attention kernels which require strictly contiguous KV cache tensors. However, when KV transfer (P/D disaggregation) is enabled via NixlConnector, the KV cache may be permuted into non-contiguous views, causing an AssertionError in is_strictly_contiguous(kv_cache_permute).

This fix automatically disables TRTLLM attention for both prefill and decode when kv_transfer_config is set, falling back to FlashInfer native kernels which handle non-contiguous tensors correctly.

Test Plan

Manually tested on DGX B200 with NixlConnector:

CUDA_VISIBLE_DEVICES=0 \
VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \
vllm serve Qwen/Qwen3-0.6B \
  --port 8100 \
  --enforce-eager \
  --block-size 64 \
  --gpu-memory-utilization 0.3 \
  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both"}'

Test Result

Before fix: Server crashes during startup with:

AssertionError: assert is_strictly_contiguous(kv_cache_permute)

After fix: Server starts successfully with info message:

INFO [flashinfer.py:582] TRTLLM attention is disabled because KV transfer (P/D disaggregation) is enabled. TRTLLM attention requires strictly contiguous KV cache tensors which may not be guaranteed with KV transfer.

Verified endpoint works:

curl http://localhost:8100/v1/completions \
  -H "Content-Type: application/json" \
  -d '{"model": "Qwen/Qwen3-0.6B", "prompt": "Hello", "max_tokens": 5}'
# Returns: {"choices":[{"text":" *everyone*, I have",...}]}

---

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 a critical bug where vLLM crashes on Blackwell GPUs when using NixlConnector for P/D disaggregation due to FlashInfer auto-enabling TRTLLM attention, which requires contiguous KV cache tensors. The fix disables TRTLLM attention when kv_transfer_config is set, falling back to FlashInfer native kernels. The changes involve modifying vllm/v1/attention/backends/flashinfer.py to implement this logic. The review focuses on the correctness of the fix and the clarity of the code.

[gemini-code-assist]

This block of code disables TRTLLM attention when KV transfer is enabled. It's crucial to ensure that this disabling mechanism is robust and doesn't inadvertently affect other scenarios where TRTLLM attention could be beneficial. Consider adding a more specific check to ensure TRTLLM is only disabled when the non-contiguous KV cache issue is present, perhaps by checking the specific KV connector being used.

[gemini-code-assist]

This check disables TRTLLM attention for prefill when KV transfer is enabled. It's important to verify that this disabling mechanism is correctly applied and doesn't inadvertently affect other scenarios where TRTLLM attention could be beneficial. Ensure that this condition is specific to the KV transfer configuration that causes the non-contiguous KV cache issue.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{Bugbot Autofix is OFF. To automatically fix reported issues with Cloud Agents, enable Autofix in the Cursor dashboard.}

Comment @cursor review or bugbot run to trigger another review on this PR

[cursor]

get_cudagraph_support ignores KV transfer disabling TRTLLM

Medium Severity

The get_cudagraph_support() classmethod returns AttentionCGSupport.UNIFORM_BATCH when TRTLLM is hardware-supported, but doesn't account for KV transfer being enabled. With the new fix disabling TRTLLM when kv_transfer_config is set, this method will incorrectly report UNIFORM_BATCH support when the actual runtime support is only UNIFORM_SINGLE_TOKEN_DECODE. This causes incorrect cudagraph mode selection in gpu_model_runner.py, as the system believes it has TRTLLM-backed uniform batch support when it doesn't.

 

[robertgshaw2-redhat]

cc @NickLucche

[NickLucche]

Thanks for spotting this @ZhanqiuHu !

@robertgshaw2-redhat @mgoin how critical is the TRTLLM kernel performance improvement?
I believe we can ensure contiguousness with this suggestion https://github.com/vllm-project/vllm/pull/33192/changes#r2736761001 (so defaulting to VLLM_KV_CACHE_LAYOUT=NHD in this case) at the cost of dropping heterogeneous TP support.
Also, do MLA models go down a different selection path for backend?

Bottom line, I agree with the proposed conservative approach to disable it altogether, but would like to flesh out implications.

[NickLucche]

we should probably use get_kv_cache_layout to check layout is != NHD which would ensure logical<>physical layout match, but also disable heterogeneous TP deployments.

@ZhanqiuHu Can you verify this is the case, so that if one starts an instance on b200 with
VLLM_KV_CACHE_LAYOUT=NHD it doesnt crash?

[ZhanqiuHu]

Hi @NickLucche, I tested using get_kv_cache_layout() != "NHD" instead of checking kv_transfer_config.

Although, on B200, even with VLLM_KV_CACHE_LAYOUT=NHD, the layout is still forced to HND (with _KV_CACHE_LAYOUT_OVERRIDE), and it triggers the check to disable TRTLLM:

Notes:

• get_kv_cache_layout() in utils.py - override takes priority over env var
• selector.py setting override - calls set_kv_cache_layout() to set _KV_CACHE_LAYOUT_OVERRIDE
• FlashInfer.get_required_kv_cache_layout() - returns HND on B200

[NickLucche]

okay so FI is forcing HND. Let's go with your original approach to stay on the safe side