Repository navigation
[M8 1/3] Extract attention contracts and add automatic backend selection - #46
Merged
Merged
Conversation
BuckMulligan99
added a commit
to BuckMulligan99/vllm-rlt
that referenced
this pull request
Sep 24, 2026
…on views KVCacheManager combines logical allocation, device storage, prepared-batch metadata and attention dispatch, and the scheduler, preemption manager, PD worker and model runner reached into its private state: _refs to price prefix hits, _get_allocation for held pages, leases and block tables, raw key_cache page copies for preemption, and duplicated written-state bookkeeping for the asynchronous finalize kernel. Add docs/kv_cache_walkthrough.md with the pinned baseline, caller inventory, allocation lifecycle, state ownership, function contracts, call chains, findings, the migration plan aligned with ThinkFlowLab#46 and ThinkFlowLab#50, and the M3 refactoring-category checklist applied to M4. Add allocation views (has_allocation, allocated_blocks, plane_block_tables, has_transfer_lease, exclusive_prefix_blocks), a written-state contract (token_written, mark_finalized) shared with finalize_token, and KVSnapshot with snapshot/restore for preemption. Migrate the M3 and M9 callers and ModelRunner.finalize_many to them. One explicit behavioral fix (F6): _claim evicted prefix-cache entries before discovering that the claim still could not be satisfied, so a failed growth, resumption or PD reservation discarded reusable prefixes. It now refuses before evicting when num_free_blocks is insufficient; successful claims are unchanged. Covered by a regression test that fails on the old code. Validation is stricter in restore, snapshot and the finalize precondition, none reachable through existing callers. The remaining private callers are the M5 prepared-batch builders, the model and the speculative runner, deferred to the next increment. Refs ThinkFlowLab#32. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: BuckMulligan99 <65453204+BuckMulligan99@users.noreply.github.com>
BuckMulligan99
added a commit
to BuckMulligan99/vllm-rlt
that referenced
this pull request
Sep 24, 2026
…on views KVCacheManager combines logical allocation, device storage, prepared-batch metadata and attention dispatch, and the scheduler, preemption manager, PD worker and model runner reached into its private state: _refs to price prefix hits, _get_allocation for held pages, leases and block tables, raw key_cache page copies for preemption, and duplicated written-state bookkeeping for the asynchronous finalize kernel. Add docs/kv_cache_walkthrough.md with the pinned baseline, caller inventory, allocation lifecycle, state ownership, function contracts, call chains, findings, the migration plan aligned with ThinkFlowLab#46 and ThinkFlowLab#50, and the M3 refactoring-category checklist applied to M4. Add allocation views (has_allocation, allocated_blocks, plane_block_tables, has_transfer_lease, exclusive_prefix_blocks), a written-state contract (token_written, mark_finalized) shared with finalize_token, and KVSnapshot with snapshot/restore for preemption. Migrate the M3 and M9 callers and ModelRunner.finalize_many to them. One explicit behavioral fix (F6): _claim evicted prefix-cache entries before discovering that the claim still could not be satisfied, so a failed growth, resumption or PD reservation discarded reusable prefixes. It now refuses before evicting when num_free_blocks is insufficient; successful claims are unchanged. Covered by a regression test that fails on the old code. Validation is stricter in restore, snapshot and the finalize precondition, none reachable through existing callers. The remaining private callers are the M5 prepared-batch builders, the model and the speculative runner, deferred to the next increment. Refs ThinkFlowLab#32. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: BuckMulligan99 <65453204+BuckMulligan99@users.noreply.github.com>
1 task done
Collaborator
|
Please remove benchmarks dir and walkthrough.md |
bjf-frz
reviewed
Sep 30, 2026
|
|
||
| def create_backend( | ||
| backend: str, device: torch.device, dtype: torch.dtype, head_dim: int, block_size: int | ||
| ): |
Collaborator
There was a problem hiding this comment.
Could we add an auto selection mode here that considers the current hardware, installed backend packages, and configuration constraints such as dtype, head dimension, and block size? Explicit backend selection should remain supported and fail clearly when incompatible. Please also report the selected backend and the selection reason.
|
|
||
|
|
||
| def backend_capabilities(implementation) -> BackendCapabilities: | ||
| return BackendCapabilities(packed_prefill=getattr(implementation, "generation", None) == 4) |
Collaborator
There was a problem hiding this comment.
Could each backend or its adapter explicitly declare its capabilities instead of inferring packed_prefill from generation == 4? This still couples the capability interface to FlashAttention versioning and does not accommodate other implementations that support packed prefill. Callers should route through the declared capabilities.
Levius-Fubuki
changed the title
Contributor
Author
|
@bjf-frz cc，thanks |
1 task done
Collaborator
|
Contributor
Author
|
Resolved the merge conflicts and moved the CLI documentation to |
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
2 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.
Backend dispatch and synchronous attention metadata currently live inside KVCacheManager, and the first M8 increment inferred packed-prefill support from the FlashAttention generation. This PR extracts those contracts into
vllm_rlt/attentionand adds an opt-inautobackend selection mode for local inference and serving.Auto tries compatible FA4, FA3, FA2, Triton, then Torch. It checks hardware, dtype, head dimension, page size and actual implementation imports. For example, SM 8.9 with 16-token pages skips FA2's 256-token page requirement and selects Triton if installed. Explicit choices keep their validation/import failures. Selection diagnostics include the requested backend, selected implementation and reason in the startup log and
attention_info.Each kernel adapter declares its packed-prefill, CUDA Graph and GPU async capabilities. The engine selects a compatible adapter before memory profiling and passes that adapter to KVCacheManager. Execution requirements filter auto candidates, so graph/async/UVA configurations fail at startup if no compatible backend exists.
The host metadata planner owns page-table padding, causal context lengths and packed query grouping. It integrates with the current multi-depth preparation path: depth-independent views remain shared, and metadata is staged in two pinned host tensors with nonblocking copies. M4 retains allocation/depth/prefix validity; M5 retains static, async, UVA and graph buffer/stream lifetimes. The speculative runner routes packed prefill through declared capabilities. CLI changes are documented in the Command-line inference section of
docs/launching.md; this PR has no main README changes. The new benchmark scripts and walkthrough have been removed.Validation of the current revision on local CPU:
--attention-backend autogenerated outputs and reported Torch selection.Additional regressions cover host-only metadata planning, two staging calls across multiple depth assignments, shared metadata views and capability-based speculative routing. GPU execution, kernel numerics and performance have not been revalidated on this revision.
Integrated upstream main:
24c433bc2ea3b0fd5b9f266571300a9386caa5eb. Refs #32.