[TRTLLM-14138][perf] Make FlashInfer decode plans sync-free with host-built page tables - #16073
Merged
yuxianq merged 3 commits intoJul 13, 2026
Merged
Conversation
Contributor
📝 WalkthroughWalkthroughChanges rework FlashInfer attention metadata preparation and decode planning to move computations from device-synchronous operations to host-side, pinned-memory-staged buffers with single async H2D transfers. A new flat batched index-building helper is added to KVCacheManagerV2 to support vectorized host-side decode block-table construction. ChangesFlashInfer decode planning sync reduction
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Prepare as FlashInferAttentionMetadata.prepare()
participant KVManager as KVCacheManagerV2
participant HostStage as Host staging (_staging_int32)
participant Device as GPU device buffers
participant DecodePlan as decode_plan()/decode_wrapper.plan()
Prepare->>KVManager: get_batch_cache_indices_flat(request_ids, num_blocks)
KVManager-->>Prepare: pinned int32 tensor of scaled indices
Prepare->>HostStage: stage cached token lengths, KV lengths, indptr on host
HostStage->>Device: async H2D copy into stable buffers
Prepare->>Prepare: retain _host_pool_indices, _paged_kv_indptr_decode_host
DecodePlan->>Prepare: _build_decode_block_tables()
Prepare->>Device: gather pool buffer rows, write to persistent block-table buffer
Prepare-->>DecodePlan: block_tables (or None to trigger fallback)
DecodePlan->>DecodePlan: decode_wrapper.plan(block_tables=block_tables)
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tensorrt_llm/_torch/attention_backend/flashinfer.py (2)
1242-1244: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winStage
decode_blocksthrough pinned memory before moving to CUDA.This CUDA-graph refresh path still builds
nbfrom a pageable CPU tensor, sonon_blocking=Truemay not avoid a host-side stall. Consider adding an int64 staging helper or using a pinned CPU tensor here.Suggested fix
- nb = torch.tensor(decode_blocks[:rows], - dtype=torch.int64).to(device=bt.device, - non_blocking=True) + nb_host = maybe_pin_memory( + torch.tensor(decode_blocks[:rows], dtype=torch.int64)) + nb = nb_host.to(device=bt.device, non_blocking=True)🤖 Prompt for 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. In `@tensorrt_llm/_torch/attention_backend/flashinfer.py` around lines 1242 - 1244, The CUDA-graph refresh path in flashinfer.py builds nb from a pageable CPU tensor, so the non_blocking transfer may still stall on the host. Update the decode_blocks handling in the refresh logic to stage the int64 data through pinned memory before moving it to bt.device, either by reusing an existing pinned staging helper or by creating the CPU tensor with pinned memory first. Keep the fix localized around the nb construction so the CUDA transfer remains truly asynchronous.
876-883: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winPin the fallback page-index tensor before async H2D copies.
The flat path returns a staged/pinned CPU tensor, but the fallback creates a pageable
torch.tensor; the latercopy_(..., non_blocking=True)can still serialize H2D for managers withoutget_batch_cache_indices_flat().Suggested fix
- return torch.tensor(indices_list, dtype=torch.int32) + return maybe_pin_memory(torch.tensor(indices_list, dtype=torch.int32))🤖 Prompt for 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. In `@tensorrt_llm/_torch/attention_backend/flashinfer.py` around lines 876 - 883, The fallback path in flashinfer attention backend returns a pageable page-index tensor, unlike the flat path that stages pinned CPU memory, so the later non_blocking H2D copy may serialize. In the fallback branch that builds the tensor from get_batch_cache_indices, make sure the returned tensor is pinned before it is used for async copy, matching the behavior of the flat-path code paths in flashinfer.py.
🤖 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 `@tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py`:
- Around line 2860-2878: The batched page-index transform in the cache-index
path is doing arithmetic in int32 via the temporary buffer from the request-ids
loop, so `out * scale // div_factor` can overflow before `BAD_PAGE_INDEX` is
masked out. Update the logic in `kv_cache_manager_v2.py` to perform the
multiply/divide in int64 for the `out` array (after it is filled from
`get_base_page_indices` and before `np.copyto`), then cast the results back to
the original int32 buffer so the batched behavior matches the old per-element
Python loop.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/flashinfer.py`:
- Around line 1242-1244: The CUDA-graph refresh path in flashinfer.py builds nb
from a pageable CPU tensor, so the non_blocking transfer may still stall on the
host. Update the decode_blocks handling in the refresh logic to stage the int64
data through pinned memory before moving it to bt.device, either by reusing an
existing pinned staging helper or by creating the CPU tensor with pinned memory
first. Keep the fix localized around the nb construction so the CUDA transfer
remains truly asynchronous.
- Around line 876-883: The fallback path in flashinfer attention backend returns
a pageable page-index tensor, unlike the flat path that stages pinned CPU
memory, so the later non_blocking H2D copy may serialize. In the fallback branch
that builds the tensor from get_batch_cache_indices, make sure the returned
tensor is pinned before it is used for async copy, matching the behavior of the
flat-path code paths in flashinfer.py.
🪄 Autofix (Beta)
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: da1effb6-0dd0-44b3-8e32-654516df56a6
📒 Files selected for processing (2)
tensorrt_llm/_torch/attention_backend/flashinfer.pytensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
This was referenced Jul 7, 2026
kaiyux
changed the title
Member
Author
|
/bot run |
Collaborator
|
PR_Github #58127 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58127 [ run ] completed with state
|
kaiyux
force-pushed
the
perf/flashinfer-sync-free-plan-20260707
branch
from
July 8, 2026 05:45
86bbdb8 to
df4bb35
Compare
Member
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #58163 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58163 [ run ] completed with state
|
Member
Author
|
/bot run |
Collaborator
|
PR_Github #58334 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58334 [ run ] completed with state |
Collaborator
|
[by Codex] @yizhang-nv Could you review this PR? Thanks! |
yuxianq
reviewed
Jul 9, 2026
yuxianq
reviewed
Jul 9, 2026
yuxianq
reviewed
Jul 9, 2026
yuxianq
reviewed
Jul 9, 2026
yuxianq
reviewed
Jul 9, 2026
yuxianq
reviewed
Jul 9, 2026
yuxianq
reviewed
Jul 9, 2026
yuxianq
reviewed
Jul 9, 2026
yuxianq
reviewed
Jul 9, 2026
yuxianq
reviewed
Jul 9, 2026
yuxianq
reviewed
Jul 9, 2026
yuxianq
reviewed
Jul 9, 2026
yuxianq
reviewed
Jul 9, 2026
yuxianq
reviewed
Jul 9, 2026
yuxianq
reviewed
Jul 9, 2026
…age tables Remove per-step GPU arithmetic, stream syncs, and scalar D2H copies from the FlashInfer attention metadata preparation path: - Build all page-table metadata (kv lens, block counts, decode/prefill/ mixed indptrs, last-page lens) with numpy on the host and stage it through pinned tensors for async H2D copies. - Add KVCacheManagerV2.get_batch_cache_indices_flat(): gather only the live block-table widths straight into a pinned CPU tensor with one vectorized transform instead of materializing padded per-request Python lists. get_batch_cache_indices keeps its trimmed-list semantics for non-flat callers. - Retain host copies of the decode indptr and per-pool flat page indices, and pre-build persistent trtllm-gen decode block tables on the host, so BatchDecodeWithPagedKVCacheWrapper.plan() no longer triggers flashinfer's per-request rebuild loop (one stream sync plus one scalar D2H per generation request per plan). Buffers are allocated once at capacity under CUDA graphs so captured decode kernels keep reading stable addresses. - Vectorize the VSWA decode block-table refresh (one gather + where instead of a per-request slice-copy loop). Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
…plans - Store the persistent trtllm-gen decode block tables (device + pinned host) on FlashInferWrappers instead of two PlanParams-keyed dicts; _plan_with_params now creates the container once per plan_params and mutates it in place across replans so the device buffer address stays stable for CUDA graphs. - Replace the per-pool host-indices lookup in _build_decode_block_tables with a _host_paged_kv_indices mirror kept in lockstep with _paged_kv_indices at the prepare fill, VSWA swap, and restore sites. - Add KVCacheManager.get_batch_cache_indices_flat (V1) matching the V2 interface and drop the getattr feature-detection fallback. - Replace the dead indptr fallback in decode_plan with an assert; access decode_wrapper._backend directly. - Hoist the max_n capacity guard and flatten the CUDA-graph no-realloc guard in _build_decode_block_tables. - Naming: _host_paged_kv_indptr_decode, _max_num_blocks, block_tables/ host_block_tables, and full variable names in the VSWA block-table refresh. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
kaiyux
force-pushed
the
perf/flashinfer-sync-free-plan-20260707
branch
from
July 10, 2026 07:50
df4bb35 to
b5be365
Compare
Rename _staging_int32 to _to_int32_tensor and make it an inner function of FlashInferAttentionMetadata.prepare, its only user, per review feedback. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
Member
Author
|
@yuxianq Thanks a lot for the very detailed review! All the suggestions should've been addressed, and I've replied some of the questions, can you help take a look again when having a chance? |
Member
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #58656 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58656 [ run ] completed with state
|
Member
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #58727 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58727 [ run ] completed with state |
yuxianq
approved these changes
Jul 13, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
FlashInfer attention metadata preparation performed per-step GPU arithmetic and synchronizations, and
BatchDecodeWithPagedKVCacheWrapper.plan()triggered flashinfer's per-request host rebuild loop — one stream sync plus one scalar D2H copy per generation request per plan. This PR makes the decode planning path sync-free by building everything on the host:prepare().KVCacheManagerV2.get_batch_cache_indices_flat(): gathers only the live block-table widths straight into a pinned CPU tensor with one vectorized transform, instead of materializing padded per-request Python lists.get_batch_cache_indiceskeeps its trimmed-list semantics for non-flat callers.plan()no longer triggers flashinfer's per-request rebuild loop. Buffers are allocated once at capacity under CUDA graphs so captured decode kernels keep reading stable addresses.No kernel behavior change: the planning inputs are identical to before, they are just produced on the host without device round-trips.
Split out of #16033 (2 of 3, together with #16072 and #16074); this PR is independent of the other two.
Test Coverage
_torchattention and executor test suites (FlashInfer backend, VSWA, and KV-cache manager coverage).045705139d) on 1x B200 (Gemma-4-31B-IT-NVFP4, ISL 1024 / OSL 128, 1024 prompts, concurrency 1024,trtllm-serve+benchmark_serving):045705139d)Together with #16072 and #16074, the three PRs compose to the +30.9% measured in #16033.
PR Checklist
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.🤖 Generated with Claude Code