[NVBUG: 6562021] Fix vLLM FlashAttention KV cache layout handling - #2084
[NVBUG: 6562021] Fix vLLM FlashAttention KV cache layout handling#2084sychen52 wants to merge 1 commit into
Conversation
📝 WalkthroughWalkthroughChangesFlashAttention detects vLLM KV-cache layouts for kv-first, blocks-first, and packed formats. Runtime extraction and GPU tests now use the detected layout. FlashAttention KV-cache layout
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant FlashAttentionForward
participant FlashAttentionBackend
participant KVCacheLayoutDetector
FlashAttentionForward->>FlashAttentionBackend: obtain KV-cache shape
FlashAttentionBackend->>KVCacheLayoutDetector: inspect installed layout
KVCacheLayoutDetector-->>FlashAttentionBackend: return supported layout
FlashAttentionBackend-->>FlashAttentionForward: return extraction rule
FlashAttentionForward->>FlashAttentionForward: extract key/value caches
Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
The fix itself looks right: probing FlashAttentionBackend.get_kv_cache_shape(3, 16, 1, 16) with distinct num_blocks=3 disambiguates (2, num_blocks, ...) from (num_blocks, 2, ...) (a plain shape[0] == 2 check would be ambiguous when num_blocks == 2), the unbound views keep the right page_size = key_cache.shape[1] in both layouts, and non-contiguous views are already supported on the FlashInfer path.
One concrete gap: tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.py was not updated and still builds KV caches in the legacy [2, num_blocks, page, heads, dim] layout (torch.stack([k_cache, v_cache], dim=0), torch.zeros(2, 1, 16, 2, 64), and the assert kv_cache.shape == (2, ...) in test_page_size_inferred_from_k_cache). On exactly the newer vLLM this PR targets, those tests will hit the same unpack error the PR fixes (or silently slice the wrong axis), so they'll go red on the new layout. That file's module docstring also still documents kv_cache.unbind(0).
Secondary points below (layout-detection logic now duplicated in tests; unused backend_cls parameter; forward-path coverage of the new layout depends on the installed vLLM version). No licensing or design-review concerns.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py`:
- Around line 50-58: Update _flash_attention_kv_cache_split_dim and the KV-cache
splitting logic before the existing unbind() call to recognize and validate the
complete supported vLLM layouts: the 5-D v0.20.0 cache and packed 4-D v0.26.0
cache. Reject malformed or unsupported backend shapes rather than relying on the
leading dimensions, and extract K/V using the layout-specific logic for the
detected format.
In `@tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py`:
- Around line 544-551: The _flash_attention_kv_cache helper duplicates
production layout-selection logic; update it to call
_flash_attention_kv_cache_split_dim and use the returned split dimension when
constructing the cache shape. Remove the probe_shape branching so layout mapping
remains centralized while this helper only builds the tensor shape.
- Around line 554-572: Extend the sparsity attention tests around
ModelOptSparseAttentionImpl.forward to parameterize both supported 5-D KV-cache
layouts and construct caches through the actual backend behavior rather than
only mocking _flash_attention_kv_cache_split_dim. Verify forward succeeds with
each layout, then add a focused case for an unsupported layout that asserts the
expected error. Reuse the existing fixtures and forward-test setup, keeping the
tests lean.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 089ce881-46aa-4611-8de4-c0f6df82c63d
📒 Files selected for processing (2)
modelopt/torch/sparsity/attention_sparsity/plugins/vllm.pytests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2084 +/- ##
===========================================
+ Coverage 67.15% 77.49% +10.33%
===========================================
Files 521 522 +1
Lines 59857 60941 +1084
===========================================
+ Hits 40199 47224 +7025
+ Misses 19658 13717 -5941
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
96e735e to
8cfd8ba
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Re-review of the KV-cache layout fix: all critical items from the previous round are now resolved.
Addressed:
- (critical)
test_vllm_plugin.pystill hardcoded the legacy layout — now builds caches through_stack_paged_cache()/vllm_plugin._flash_attention_kv_cache_split_dim(),test_chunked_prefill_is_forwarded_to_kernelno longer hardcodestorch.zeros(2, 1, 16, 2, 64), andtest_page_size_inferred_from_k_cachenow asserts on the unboundkey_cache.shapeinstead of the raw(2, ...)cache shape. Module docstring updated too.test_prefill_matches_contiguoustherefore also gives a numerical check that the blocks-first non-contiguous K/V views work on the targeted vLLM. - (critical) forward-path coverage was version-dependent — new
test_flash_attention_forward_follows_backend_kv_cache_layoutparameterizessplit_dim in {0, 1}, monkeypatchesFlashAttentionBackend.get_kv_cache_shape, clears thefunctools.cachebefore/after, and assertsk_cache/v_cachedata_ptr()pluspage_size == 16. Verified this genuinely discriminates the axes: on(3, 2, 16, H, D)the value-cache offset differs betweenunbind(0)andunbind(1), so a wrong-axis regression fails. - (minor) duplicated layout mapping in tests — both test modules now consume the plugin helper;
_flash_attention_kv_cacheonly builds the shape viashape.insert(split_dim, 2). - (minor) unused
backend_clsparameter — dropped; the helper referencesFlashAttentionBackenddirectly.
Residual, non-blocking: CodeRabbit's claim that some newer vLLM returns a packed 4-D cache (num_blocks, num_kv_heads, block_size, 2*head_size) is not handled — on such a version _flash_attention_kv_cache_split_dim() raises RuntimeError at first forward (and the GPU test helpers would error rather than skip). That's an explicit failure rather than silent wrongness and is outside the scope of the reported NVBug, but if you expect to support that release, consider widening the detector and including the full probed shape (not just the leading two dims) in the error message. No licensing or design-review concerns; the bot "autofix prompt" blocks in the PR conversation are standard CodeRabbit tooling output, not directives I acted on.
Complex PR: 2 existing test files modified or removed. Looping in a human for approval.
Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
8cfd8ba to
f99adfe
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py`:
- Around line 50-60: Update _flash_attention_kv_cache_layout and the kv_cache
unpacking flow to classify physical layout with
FlashAttentionBackend.get_kv_cache_stride_order(), not shape alone, and
normalize NHD/HND caches before dispatch. Validate runtime rank and dimensions
for each supported 4-D or 5-D layout, preserving correct page/head axis ordering
for both stride orders. Add independent tests covering NHD and HND layouts,
including the packed branch.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3e0770ba-5db5-481b-834b-cdc33003ede5
📒 Files selected for processing (3)
modelopt/torch/sparsity/attention_sparsity/plugins/vllm.pytests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.pytests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.py
- tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Re-review of the FlashAttention KV-cache layout fix (3 files, +126/-25). All critical items from the previous rounds are now resolved, and the previously-flagged 4-D gap is closed — that note is withdrawn.
Addressed since the last round:
- (critical) packed 4-D layout unhandled —
_flash_attention_kv_cache_layout()now returns"kv-first"/"blocks-first"/"packed"and the packed branch doeskv_cache.transpose(1, 2).split(self.head_size, dim=-1). Checked the shapes:(blocks, heads, page, 2*dim) -> (blocks, page, heads, dim)for both K and V, so the downstreampage_size = key_cache.shape[1]contract still holds in all three layouts, andassert kv_cache.shape[2] == page_sizeintest_page_size_inferred_from_k_cacheis valid for every layout. - (critical)
test_vllm_plugin.pyhardcoded the legacy layout — now builds caches through_make_backend_paged_cache(), which consumes the plugin's layout helper;test_chunked_prefill_is_forwarded_to_kernelno longer hardcodestorch.zeros(2, 1, 16, 2, 64); module docstring updated. - (critical) version-dependent forward coverage —
test_flash_attention_forward_follows_backend_kv_cache_layoutis now parameterized over all three layouts with a monkeypatchedget_kv_cache_shape,cache_clear()before/after, and shape/stride/data_ptr()assertions, so a wrong-axis (or K/V-swapped) regression fails regardless of the installed vLLM. - (minor) unused
backend_cls— dropped; the helper referencesFlashAttentionBackenddirectly. - (minor) duplicated layout mapping in tests — both test modules now derive the cache shape from the plugin helper.
Remaining for owner judgment (not re-flags of fixed bugs):
- Packed K/V ordering is an unverified assumption about vLLM. The packed branch assumes K is the first
head_sizeslice of the last dim and V the second. Every test constructs the cache itself with the same convention (torch.cat([k_cache, v_cache], dim=-1)), so if the release that introduced(num_blocks, num_kv_heads, block_size, 2*head_size)actually writes V first (or interleaves), the plugin would silently swap K and V — wrong numerics with no test failure, since the self-consistent fixtures never exercise vLLM's ownreshape_and_cache_flashwrite. Worth confirming against the vLLM source/an e2e accuracy run on that version before shipping. - No coverage of the
RuntimeErrorpath for an unrecognized backend shape (CodeRabbit asked for this); the message does now include the full probed shape, so this is small. - Minor probe ergonomics: the detector probes
get_kv_cache_shape(3, 16, 1, 16), i.e.block_size == head_size == 16, so the exact-tuple comparisons rely on coincidence-free dims. Distinct values (e.g.head_size=32) would make the three mappings unambiguous by construction. - Test/impl mirroring: the split expression is now written in three places (plugin, plus the expected value in two tests). A single plugin-level
_split_kv_cache(kv_cache)reused by the tests would keep the assertions from restating the implementation.
💬 Note on the conversation: CodeRabbit's get_kv_cache_stride_order() finding was withdrawn after the author's reply that only the logical layout matters — that looks right here (the Triton path consumes strided views, as the FlashInfer NHD/HND test shows), so it isn't a blocker. The CodeRabbit "autofix prompt" blocks in the PR body/comments are standard tooling output and were treated as data, not instructions.
Licensing: none. Design review: not applicable (localized compatibility fix).
2 participants
What does this PR do?
Type of change: Bug fix
The ModelOptSparseAttentionImpl.forward (FlashAttention path) assumes the paged KV cache has shape [2, num_blocks, page_size, num_kv_heads, head_dim], where dimension 0 is the K/V split. unbind(0) on this shape returns exactly 2 tensors.
Newer vLLM versions changed the FlashAttention kv_cache layout to [num_blocks, 2, page_size, num_kv_heads, head_dim] (same as FlashInfer). unbind(0) now returns num_blocks tensors, causing the unpack error.
Now, we detect the layout first and then unbind.
Usage
save as before.
Testing
unittest
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: N/AAdditional Information
Summary by CodeRabbit
Bug Fixes
Tests