Conversation
Ported from the internal branch ahead of the implementation they exercise: per-locality-domain slot pools, storage manager localization, and a stream-localization end-to-end case. They fail until the KVCacheManagerV2 localization lands. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
Two-stream attention runs each locality domain on its own stream, so the two must not share an AttentionOp. Thread locality_domain_id from the op entry point into get_attention_op and add it to the cache key. Defaults to -1, the non-localized sentinel, so existing callers are unaffected. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
Thread locality_domain_id from the request queue onto LlmRequest, and let draft_kv_cache_context opt the draft sub-forward into fork-join: it either disables locality domain for the draft, or swaps in the draft cache's per-domain block offsets when that cache supports fork-join. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
MTP and Eagle3 request locality domain fork-join for the draft sub-forward, and the drafter carries the parent request's locality domain onto the draft request. The flag stays False by default, so other callers keep running the draft single-stream. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
Split the VMM allocator and GPU slot pool from one physical backing to one per locality domain behind a single virtual reservation, so KV pages live in the half of the die that will read them. Both keep their single-pool construction path, defaulting to domain 0, so existing call sites are unchanged. Carries forward two main-side changes the internal branch had not picked up: the POSIX FD handle-type preference from NVIDIA#18347, and the SlotId annotation on GpuSlotPool.slot_address from NVIDIA#16218. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
…tocol Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
Thread locality_domain_id through slot allocation, free-slot preparation, and page indexing. Callers move to keyword arguments so the new parameter cannot bind positionally. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
Reject enable_locality_domains on the C++ backend instead of discarding it, and skip the localization tests unless the Python backend is selected. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
Port the page_filter parameter the storage manager already calls, fix a stale KVCacheManagerV2 import path, and skip the storage-manager tests unless the Python backend is selected. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
TestKVCachePlacement and TestLocalizedInvalidValueCheck exercise KVCacheManagerV2, which does not carry locality domains yet. They return with that layer. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
Adds the ResourceManager.pick_locality_domain hook, defaulting to no domain, and threads its result through request construction. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
Introduces the LlmArgs flag that turns locality domain execution on, the matching golden-manifest entry, and the post-transform feature the weight loader qualifies against. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
Passes the locality domain policy into ModelConfig, publishes it for the weight loader, enables the localized-weights post-transform feature, and rejects the GMS load format, which shares registered parameters that locality domain execution releases. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
The flag reaches _create_kv_cache_manager for both the main and draft models. Forwarding it to KVCacheManagerV2 lands with that class. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
Adds the num_locality_domains query, the default request-id placement policy, request-id validation, backend-compat constructors for GpuCacheTierConfig and ReuseScope, and the cached fork-join flag. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
Requests and dummy requests carry a domain into cache creation, validated against the manager's domain count. The impl call only forwards the id to backends whose signature accepts it, so the C++ manager is unaffected. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
A localized GPU pool spans one virtual address region per domain with an unmapped gap between them, so a single flattened view would read the gap. Return one contiguous view per domain instead. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
get_locality_domain reports where an allocated request landed and get_per_locality_domain_free_slots reports per-domain capacity, with a fallback for backends that do not expose it. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
…ackend Declares the per-domain metadata fields the fork-join path needs and pre-allocates one attention workspace per domain. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
Each iteration derives per-domain request and token boundaries from the KV cache manager's placement, then packs matching per-domain block-offset tensors so each domain's attention kernel sees only its own requests. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
Each domain gets its own host launch scalars and packed block offsets. Under CUDA graphs the scalars use conservative bounds so a graph captured at one decode length replays at longer ones with the same request split. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
Splits per-iteration attention inputs into per-domain views: token boundaries for generation-only batches, sliced cumulative sequence lengths, and per-domain MLA generation state derived from the full batch. Signed-off-by: Chulian Zhang <851104+zhangcl@users.noreply.github.com>
Sort context and generation requests by locality domain, publish the per-domain request/token split points on attn_metadata, and give each active domain one max-length dummy request during CUDA graph warmup.
The KV cache manager owns the placement policy; non-V2 managers inherit the base no-op and keep returning None.
A single thread-local wrapper is wrong once each locality domain runs under its own green context.
A/B against the pre-change engine shows the split and token budget are unchanged, so pin the ids and the measured budget.
base_address is the address of slot 0, which lives in domain 0, so the slot count has to describe the same domain.
Host and disk storage inherit the base signature, so spilling a domain-pinned cache raised TypeError. Gate on localized_gpu_storage, which both callers already compute.
shrink/expand acted on locality domain 0 only, overflow was classified by canonical slot id against a domain-local capacity, defragmentation drew destination slots from domain 0, and the resize path never divided the aggregate quota the way construction does.
pad_batch trims padding dummies positionally with generation_requests[:-padding_size], but the locality domain sort could move them off the tail and drop real requests instead.
self.head_dim is the raw constructor argument and may be a list, which breaks the view shape for per-layer head_dim models such as Gemma4 hybrid.
Drop getattr guards on states that cannot occur, the redundant extra_attrs sync, and the hand-rolled context re-binning; gate request placement on _is_kv_manager_v2 so BaseResourceManager no longer needs a no-op override.
batched_lock_to_gpu reserved GPU slots in a request's locality domain but let _batched_migrate allocate the destination from domain zero, silently placing resumed pages in the wrong half of the GPU. Pass the domain there, and assert it is supplied whenever the destination pool group is localized so a missed call site fails instead of misplacing pages. The spill and prefetch call sites migrate away from the GPU tier today, so passing the domain there is defensive only. The docstring claimed a cross-tier destination is always Host/Disk, which is what made the domain-zero default look safe; it is not true for the load-back direction.
The locality domain metadata was prepared but never consumed: nothing called _build_locality_domain_slices, nothing entered LocalityDomainRuntime, and the attention op always received locality_domain_id=-1. _attention_partitions yields the batch unchanged when locality domains are off and one partition per domain when they are on, each with its rows, block offsets, workspace and MLA generation state in place and its stream current, so forward keeps a single dispatch site for both paths. Generation-only for now; context and mixed batches take the unpartitioned path. Verified on Rubin sm_107 with DeepSeek-V3-Lite: 1170 per-domain dispatches across two domains, output identical to the non-localized run.
mla_rope_generation stays a single full-batch op, but the cu_seqlens and FMHA scheduler counter it produces are per-request. Without splitting them each domain received the whole batch's values and both shared one scheduler counter, which the two streams increment concurrently. Verified on Rubin sm_107 with DeepSeek-V3-Lite and real (non-mocked) localization: output identical to the non-localized run.
Extends _attention_partitions beyond generation-only, so context attention also
runs per locality domain instead of falling back to a single kernel over the
whole batch.
Prep for every domain completes before runtime.fork(). fork() records an event
on the current stream and has the partition streams wait on it, so anything
enqueued afterwards is unordered with respect to them: building a domain's
cu_seqlens after the fork let its attention kernel read them before the
host-to-device copy and cumsum had run, and those values index q/k/v, so the
kernel walked out of bounds. Generation-only never hit this because its slices
are views with no producer kernel.
Also, from MR !10223:
- k/v gathered with the dense-KV range (ctx_kv_token_*), not the new-token
range used for q / q_pe / latent_cache
- total_gen_tokens passed only for generation_only
- per-domain MLA gen state applied only for generation_only; context/mixed
rebuild cu_q_seqlens, cu_kv_seqlens, a scheduler counter, quant_q_buffer
- workspace growth on domain 0 propagated to the other domains
- block_ids_per_seq is indexed over the whole batch, so pass None per domain
- scatter and workspace resize stay inside partition_context
Exercised end to end on a localization-capable device with real (non-mocked)
localization, both with and without CUDA graphs, over repeated runs with the
partition streams concurrent; output matches the non-localized baseline in each
configuration.
…FP4 output TRTLLM_LOCALITY_DOMAIN_TWO_STREAM_ATTN=0 forks and joins around each locality domain in turn instead of once around all of them, so the domains run one at a time while each keeps its own stream, workspace and metadata. That separates a partitioning bug from a race between the domains, which CUDA_LAUNCH_BLOCKING cannot do because it serializes everything. NVFP4 output quantization now asserts instead of silently running the batch unpartitioned, matching MR !10223. Behaviour change: a configuration that combined output_sf with locality domains used to degrade quietly and now fails. Verified by counting fork/join: concurrent does 1230 of each, serialized 1920, and the 690 difference equals the number of localized attention calls, the remainder coming from the autotuner's own fork. Output is identical in both modes.
… manager MambaHybridCacheManagerV2 overrides _create_kv_cache, so adding locality_domain_id to the base signature broke it: both call sites in KVCacheManagerV2 pass the argument unconditionally, without checking num_locality_domains, so every hybrid model on KV cache v2 failed at request creation with "got an unexpected keyword argument" even with locality domains disabled. The override exists only on main -- on the branch this was ported from the subclass inherits the base directly -- so the conflict is a product of the port rather than of the original change. The recurrent state itself is not locality-domain placed, but the hybrid manager also owns the KV cache for the model's attention layers, so the argument is forwarded rather than dropped. Found by running a hybrid model end to end; the failure is gone afterwards.
…ion test The localized arm only flipped the mock capability variable, so both arms constructed an ordinary LLM and the equivalence assertion held no matter what locality domains did. Pass enable_locality_domains and select the Python KV cache backend, which is the only one that implements them. Instrumenting LocalityDomainRuntime confirms the difference: the test now drives thousands of paired fork/join calls where it previously drove none.
Review feedback: the two add_dummy_requests calls differ only in arguments that already collapse to the same values. With one long request per locality domain and none configured, num_long_requests is 1, so max_seq_len_request_ids is [batch_size - 1] and the keyword dict is empty -- exactly what the mixed-context branch passed explicitly. Keeping them apart also dropped locality_domain_ids on the mixed-context path, since only the other branch forwarded the keyword dict.
…mock Locality domains need hardware that exposes two locality domains, so the test is marked skip_no_rubin rather than pretending elsewhere. That also removes the reason to mock. TRT_LLM_MOCK_LOCALIZATION_SUPPORT only reaches the KV cache allocator: the fork/join runtime asks the localization handle directly, so enabling locality domains under the mock fails on hardware that does not support them. On Rubin the capability is real, and probing the device exercises real pools, streams and compute split instead of a stand-in. Verified on Blackwell: both cases report "only supported in Rubin architecture" instead of passing without exercising the feature.
…index bound Review feedback: get_page_index_upper_bound multiplies the base term by attr.expansion but not the locality displacement, while PageIndexConverter emits index * expansion + i for every converted index. A page displaced into locality domain 1 therefore converts to an index carrying the displacement times expansion, which the bound did not reserve, so tensor views wrapping the pool come out undersized. Only reachable with heterogeneous tokens_per_block: expansion is exact_div(tokens_per_block, tokens_per_block_override) and no model config sets the override today, so every current model has expansion 1 and the term is unchanged. The new test sets the override explicitly and fails without the fix.
…ong reason
Review feedback, all cases where the assertion was weaker than the behaviour
the test names:
- the page index upper bound test only required a positive value, so an
implementation that divided by page_stride still passed; it now pins the
bound to the unpatched value, which is the documented contract
- three guard tests accepted any AssertionError raised anywhere inside the
call, including one raised before the locality_domain_id check; they now
match the production message
- the quota test only required the total to be under the budget, which
allocating nothing also satisfies; it now requires every domain to get a
real and equal share
…troller
Review feedback:
- no test pinned enable_locality_domains, so neither the False default nor
the field surviving from_yaml was covered
- the localized eviction filter test left a page scheduled when the
controller was collected; under TLLM_DEBUG_MODE=1 its __del__ calls
assert_critical and terminates the process. The remaining page is now
drained in a finally block so it happens even when an assertion fails.
…ttention # Conflicts: # tensorrt_llm/_torch/attention/backends/trtllm.py
…rence references/llm.yaml is compared with assert_equal against the introspected LLM.__init__ signature, so a new field that is absent from it fails tests/unittest/api_stability. Verified: 3 failed without the entry, 64 passed with it. The field is prototype, so it is not added to references_committed.
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
yuxianq
added
the
api-compatible
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.
@coderabbitai summary
Description
KV cache locality currently depends on a fixed 1 TiB virtual-address gap and domain-encoded slot IDs. Make
Slot.locality_domain_idauthoritative and use ordinary global slot IDs, so locality-unaware consumers address ld1 pages with the same base-plus-slot-stride arithmetic.Allocate granularity-aligned slot ranges across every pool while preserving live IDs during growth and shrink. Update migration and release, allocation rollback, page-index bounds, validation views, and disaggregation memory registration to respect slot ownership and mapped ranges.
Depends on #18919, based on its commit
402f6e46d70d2a527cf3a8285c3a5dbd29f83621. The diff againstmaincurrently includes that unmerged dependency. The slot-metadata change is isolated in one commit.Test Coverage
test_storage_manager.py: 35 passed.test_kv_cache_localization.py: 24 passed.test_kv_cache_manager_v2.py: 18 passed.tests/unittest/disaggregated/region/test_page.py: 19 CPU tests passed, using isolated package imports.Locality tests use
TRT_LLM_MOCK_LOCALIZATION_SUPPORT=1. They validate metadata, address calculation, real GPU virtual-memory mappings and copies, migration, and resize behavior. Physical ld0/ld1 memory placement is not validated because the CUDA 13.2 test environment lacks the newer locality allocation API.PR Checklist
Please review the following before submitting your PR:
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.