[None][fix] Select FlashInfer hybrid default KV layer - #17382
[None][fix] Select FlashInfer hybrid default KV layer#17382mihai-chiorean wants to merge 5 commits into
Conversation
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
|
/bot run |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (7)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (6)
WalkthroughFlashInfer metadata now selects attention-backed KV layers explicitly. Page-table preparation and VSWA pool handling use the selected layer. Tests cover hybrid and recurrent-only Mamba cache managers and Qwen3.5 FlashInfer accuracy. ChangesFlashInfer layer-aware KV handling
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant CacheManager
participant FlashInferMetadata
participant VSWA
CacheManager->>FlashInferMetadata: select attention-backed default KV layer
CacheManager->>FlashInferMetadata: prepare layer-specific page indices
FlashInferMetadata->>CacheManager: return page-table metadata
CacheManager->>VSWA: activate the default layer pool
VSWA->>CacheManager: restore active page-index state
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
eopXD
left a comment
eopXD
left a comment
There was a problem hiding this comment.
With your change, we should be good to resolve FIXME under resource_manager.py given it should be safe to remove usage of blocks_in_primary_pool here.
Your merge request should enable hybrid models to run with the FlashInfer backend. Please add e2e test coverage for this.
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@tests/integration/defs/accuracy/test_llm_api_pytorch.py`:
- Line 6222: Update test_bf16_flashinfer to include an explicit None return
annotation, then register this test in both the l0_h100.yml and
llm_function_core.txt test lists so it runs in CI and manual QA.
- Around line 6221-6233: Register TestQwen3_5_4B::test_bf16_flashinfer from
accuracy/test_llm_api_pytorch.py in all applicable test lists: l0_h100.yml,
llm_function_core.txt, and llm_function_rtx6k.txt.
🪄 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: 965c688d-5a40-4c02-9dad-b13d1658650a
📒 Files selected for processing (4)
tensorrt_llm/_torch/attention_backend/flashinfer.pytensorrt_llm/_torch/pyexecutor/resource_manager.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/unittest/_torch/executor/test_mamba_cache_manager.py
💤 Files with no reviewable changes (1)
- tensorrt_llm/_torch/pyexecutor/resource_manager.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/unittest/_torch/executor/test_mamba_cache_manager.py
- tensorrt_llm/_torch/attention_backend/flashinfer.py
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
mihai-chiorean
changed the title
|
Addressed the review in
Validation:
|
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
The core selection logic is right — keying the default page-index source off num_kv_heads_per_layer[offset] > 0 matches how the hybrid managers mark recurrent layers (CppMambaHybridCacheManager.__init__ zeroes those entries), and restricting the multi-pool layer_space scan to attention layers removes a spurious pool that could have flipped the VSWA path on for plain hybrid models. Two robustness points on the sizing path and the draft-view lookup are inline.
One behavior worth calling out beyond what the description covers: _paged_kv_indices is no longer floored by blocks_in_primary_pool. For every model it is now sized purely from get_buffers(layer).shape[0] over attention layers, which changes the allocation for non-hybrid models too — smaller in the common case, but zero if get_buffers ever returns None/is absent for the selected layer. The description reads as if the change only affects hybrid models.
For pre-merge cost: test_bf16_flashinfer adds a full GSM8K run to l0_h100. If the intent is a regression guard for the page-index selection rather than accuracy coverage of the backend, a cheaper single-config run would give the same signal.
The description is otherwise accurate against the diff, and the change is internal, so no docs/release-note gap.
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
|
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. |
|
@brnguyen2 Thanks for the detailed review. The optional robustness points are addressed in |
|
|
||
| if self.kv_cache_manager is not None: | ||
| blocks_in_primary_pool = self.kv_cache_manager.blocks_in_primary_pool | ||
| pp_attention_layer_ids = self._get_pp_attention_layer_ids() |
There was a problem hiding this comment.
Why can't we directly use kv_cache_manager.pp_layers instead of _get_pp_attention_layer_ids?
| for layer_idx in pp_attention_layer_ids: | ||
| layer_buffer = get_buffers(layer_idx) | ||
| if layer_buffer is not None: | ||
| max_num_blocks = max(max_num_blocks, |
There was a problem hiding this comment.
@yizhang-nv whether it is the simplest and correct way to calculate max_num_blocks?
eopXD
left a comment
eopXD
left a comment
There was a problem hiding this comment.
LGTM, thank you for addressing the comments.
Of course! No it just needs CI. |
6 participants
Summary
Use the first attention-backed KV-cache layer as FlashInfer metadata's default page-index source. Hybrid recurrent layers can expose placeholder page IDs, which are invalid for paged attention. The change also excludes recurrent-only layers from page-table sizing and multi-pool mapping.
Page-index buffer sizing now derives from attention-layer buffers for both hybrid and non-hybrid managers. Managers that cannot expose those buffers retain the existing
blocks_in_primary_poolsizing fallback; legitimate recurrent-only ranks keep a zero-length page-index buffer, log that paged KV metadata is disabled, and clear page-table state.Testing
pre-commit run --files tensorrt_llm/_torch/attention_backend/flashinfer.py tensorrt_llm/_torch/pyexecutor/resource_manager.py tests/unittest/_torch/executor/test_mamba_cache_manager.py tests/integration/defs/accuracy/test_llm_api_pytorch.py tests/integration/test_lists/test-db/l0_h100.yml tests/integration/test_lists/qa/llm_function_core.txt tests/integration/test_lists/qa/llm_function_rtx6k.txt5 passed, covering V1/V2 hybrid layer selection, fallback sizing, and V1/V2 recurrent-only ranks.No user-facing API change.
Dev Engineer Review
blocks_in_primary_poolFIXME was removed.QA Engineer Review
Test code changed in:
tests/unittest/_torch/executor/test_mamba_cache_manager.pytests/integration/defs/accuracy/test_llm_api_pytorch.pyAdded test coverage verifies:
The Qwen3.5-4B accuracy test is registered in:
tests/integration/test_lists/test-db/l0_h100.ymltests/integration/test_lists/qa/llm_function_core.txttests/integration/test_lists/qa/llm_function_rtx6k.txtThe changed unit tests are not registered in the integration test lists. The integration accuracy test is registered for CI and QA.
Verdict: sufficient.
Validation included focused hybrid unit tests, a passing Spark SM121 / FlashInfer 0.6.15 V1 hybrid regression, and Qwen3.6-35B-A3B NVFP4 production-path smoke testing. The V2 regression was not runnable in the live H059 container because its older runtime rejected
BlockReuseConfig(policy=...)before metadata construction. V2 remains covered by pre-merge CI and current-build validation.