[None][feat] Support the masked DSA indexer k-cache pool in the Python cache transceiver - #17283
Open
Tabrizian wants to merge 6 commits into
Open
[None][feat] Support the masked DSA indexer k-cache pool in the Python cache transceiver#17283Tabrizian wants to merge 6 commits into
Tabrizian wants to merge 6 commits into
Conversation
…n cache transceiver The Python (v2) KV-cache transceiver previously raised NotImplementedError for the per-layer masked DSA indexer k-cache pool (cross-layer indexer sharing, e.g. GLM 5.2), forcing those checkpoints onto the C++ transceiver. Teach build_page_table's indexer REPLICATED view to cover only the indexer-owning layers: one buffer entry per owning layer, mapped to its packed row via get_indexer_k_cache_pool_layer_idx, and skip the pool for a layer group with no owning layers. The Python transfer machinery already matches peers per-pool by role + global_layer_id overlap, so a masked subset transfers correctly (including PP reshard) with no further changes; the dense/unmasked layout is byte-for-byte unchanged. With Python support in place, GlmMoeDsaForCausalLM prefers the Python transceiver again like the other DeepSeek-family checkpoints, reverting the CPP override that was added only because Python lacked masked-pool support. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
Tabrizian
requested review from
2ez4bz,
bo-nv,
cascade812,
reasonsolo and
tongyuantongyu
August 5, 2026 05:20
Member
Author
|
/bot run --disable-fail-fast |
Contributor
WalkthroughThe PR adds masked DSA indexer page-table support, refreshes draft-loop token mappings, changes bounce transport admission thresholds, removes transfer-budget blocking, and adds opt-in prefix-token caching. ChangesMasked DSA cache handling
Bounce transport thresholds and admission
Prefix token cache
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant DefaultInputProcessor
participant PrefixTokenCache
participant FastTokenizer
DefaultInputProcessor->>PrefixTokenCache: encode eligible prompt
PrefixTokenCache->>FastTokenizer: tokenize uncached tail
FastTokenizer-->>PrefixTokenCache: return token IDs and offsets
PrefixTokenCache-->>DefaultInputProcessor: return assembled token IDs
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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.
🧹 Nitpick comments (2)
tests/unittest/disaggregated/test_extractor.py (1)
211-215: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate
indexer_k_cache_layer_mask.Declare the new parameter as
list[bool] | None. This preserves the helper contract and matches the documented global mask format.As per coding guidelines, “Annotate every function” and “prefer built-in generic types and
|.”🤖 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 `@tests/unittest/disaggregated/test_extractor.py` around lines 211 - 215, Update the _make_v1_dsa_manager parameter annotation for indexer_k_cache_layer_mask to list[bool] | None, preserving its existing default and behavior.Source: Coding guidelines
tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py (1)
336-337: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse declared
KVCacheManagerattributes.
KVCacheManager.__init__always setsenable_indexer_k_cacheandindexer_k_cache_local_layer_mask. Replace bothgetattrcalls with direct attribute access. This keeps the manager contract type-checkable and fails fast on an invalid manager.As per coding guidelines, “Avoid reflection when ordinary explicit code is sufficient.”
🤖 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/disaggregation/resource/kv_extractor.py` around lines 336 - 337, In the KV cache extraction logic, replace both getattr calls on kv_cache_manager with direct access to its declared enable_indexer_k_cache and indexer_k_cache_local_layer_mask attributes, preserving the existing conditional behavior and allowing invalid managers to fail fast.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py`:
- Around line 336-337: In the KV cache extraction logic, replace both getattr
calls on kv_cache_manager with direct access to its declared
enable_indexer_k_cache and indexer_k_cache_local_layer_mask attributes,
preserving the existing conditional behavior and allowing invalid managers to
fail fast.
In `@tests/unittest/disaggregated/test_extractor.py`:
- Around line 211-215: Update the _make_v1_dsa_manager parameter annotation for
indexer_k_cache_layer_mask to list[bool] | None, preserving its existing default
and behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 37d069cf-5d7c-4823-a88e-1fc952bcb392
📒 Files selected for processing (4)
tensorrt_llm/_torch/disaggregation/resource/kv_extractor.pytensorrt_llm/_torch/models/modeling_deepseekv3.pytests/unittest/disaggregated/test_extractor.pytests/unittest/llmapi/test_llm_args.py
Collaborator
|
PR_Github #63946 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #63946 [ run ] completed with state
|
1 task
Cherry-pick of NVIDIA#16925 (d2ecec4, c7a7ab0, 79cb3be) squashed into one commit. req_idx_per_token is built once per engine step in prepare_for_indices_conversion() from the target forward's seq_lens, which for an MTP generation batch are (max_draft_len + 1) tokens per request. MTPEagleWorker.forward rewrites seq_lens to one token per request before calling update_for_spec_dec(), but on_update_kv_lens() deliberately reuses the map from prepare(). Nothing rebuilds it, so from the second draft iteration onward token j resolves to request j // (max_draft_len + 1) instead of request j; only row 0 is ever correct. Both directions are affected. convert_req_index_to_global() resolves the sparse top-k indices through the wrong request's block table, and Indexer._update_k_cache() scatters the draft token's indexer K through the same slot mapping into another request's pages. The damage is confined to the MTP layer's own indexer state, so generated output is unaffected and the symptom is purely a loss of acceptance length. Because attention DP splits requests across ranks, the misrouted fraction is (B - 1) / B with B the per-rank decode batch, which is why acceptance falls monotonically with concurrency and why max_draft_len = 1 is immune (the draft loop never mutates the layout in that case). DeepseekV4TrtllmAttentionMetadata already rebuilds this map in its on_update_kv_lens() override; models on the base DSA metadata class (deepseek_v32, glm_moe_dsa) do not. Rebuild it from the current seq_lens with a device-side searchsorted, which is CUDA-graph safe and a no-op on the target forward. Also carries the PR's regression tests (test_on_update_kv_lens_rebuilds_req_idx_in_draft_loop and test_on_update_kv_lens_matches_prepare_on_target_forward) plus the MockMetadata flag defaults they need. Branch adaptation on top of the upstream PR: this branch's on_update_kv_lens() also reads in_mtp_draft_loop for the cross-step indexer top-k reuse, which MockMetadata does not set, so the two new tests would raise AttributeError before reaching the rebuild. Default it to False in the mock alongside the topk flags. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
… and the transfer admission controller Two related removals in the disaggregated executor loop. 1. _check_disagg_transfer_progress_when_idle now returns immediately. It ran two collectives every executor iteration -- _sync_disagg_gen_status_entry (WORLD-scoped) and _sync_disagg_ctx_status_entry (TP/CP) -- purely to vote on whether any rank should enter a blocking transfer-status wait. Instrumented on a GLM-5.2 disagg GEN worker, those votes were 3,995 ms of the method's 4,012 ms across 3,000 calls: 99% of its cost was the voting itself. An NVTX capture of the CTX put the WORLD allreduce at 19.4% of context wall-clock (18 calls, 211 ms mean) -- it spans every rank in the deployment, so each call waits on the slowest participant anywhere, ctx or gen. 2. _apply_disagg_transfer_admission is a passthrough. It deferred gen-init requests to bound in-flight KV transfer blocks; its FCFS deferral bursts admissions and was measured costing ~18% of GEN GPU-idle time. With (1) gone there is no transfer budget left to wait on, so admitting every fitting request is the consistent choice. Transfer completion is still reaped at the other call sites in the executor loop; only the ones inside the idle check are dropped. Both replacements are rank-uniform and contain no collectives, so ranks cannot diverge -- the deadlock the voting guarded against cannot arise once nothing here is conditional. The original bodies are left in place below the early return rather than deleted, so the behaviour being disabled stays readable. MEASURED RESULT, recorded so this is not mistaken for a win: on GLM-5.2 GB300 e2e (pareto07, 9 ctx x DEP4 + dep16 gen, c1024) this configuration measured 682.3 output tok/s/GPU against a 692.3 baseline -- 1.4% WORSE. Removing the attention-DP _can_queue allgather separately measured 682.4, within 0.02% of the same number, despite the NVTX trace billing the collective it deleted at 19.4% of ctx wall. The explanation is that these barriers EXPOSE rank skew rather than cause it: their duration is the wait for the slowest participant, so deleting one relocates the wait into the model's own collectives (the MoE all-to-all already accounts for 40-64% of ctx GPU busy time) instead of recovering it. Trace attribution alone cannot distinguish "expensive" from "merely where idleness is accounted"; only an A/B can. Carried on this branch for evaluation in combination with the other changes, not because either has been shown to help on its own. (cherry picked from commit e588a93)
In multi-turn agentic serving each turn's prompt is the previous turn's prompt plus a small delta, but the frontend re-tokenizes the whole prompt every turn. On a GLM-5.2 disaggregated context server with ~38k-token prompts, nsys attributed 47.4% of context wall-clock to the tokenize-prompt range at 43.7 ms/request. Reusing the tokenization of the longest cached prefix and tokenizing only the tail brings that to 5.49 ms/request (10.5% of wall). Correctness is the whole difficulty: splitting a string and tokenizing the tail in isolation is not generally equal to tokenizing the whole, because BPE merges can straddle the seam. The cache backs off "overlap" tokens from the split, re-tokenizes from there, and requires the first "resync" re-tokenized ids to match the cached ids over the same span; if they do not, it tokenizes the prompt in full. Verified against the GLM-5.2 tokenizer on simulated multi-turn growth plus unrelated prompts: 50/50 byte-identical token-id sequences, 39 cache hits, 0 resync fallbacks. Off by default (TLLM_PREFIX_TOKEN_CACHE=1 to enable) and applied only to plain long text prompts with a fast tokenizer -- skipped whenever the arguments would change tokenization (add_special_tokens, truncation, or a separate query). Any exception inside the cache falls back to the normal path, so it cannot fail a request. Measured end to end on GB300 as part of the ctx tuning that took GLM-5.2 pareto07 from 692.3 to 739.7 output tok/s/GPU. (cherry picked from commit a99b279)
Signed-off-by: Simeng Liu <simengl@nvidia.com>
Signed-off-by: Simeng Liu <simengl@nvidia.com>
SimengLiu-nv
requested review from
PerkzZheng,
fredricz-20070104,
ruodil and
yuxianq
August 6, 2026 23:06
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
tensorrt_llm/inputs/prefix_token_cache.py (1)
30-30: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd complete cache API contracts.
Add return types for
_bucket_key()and_find_longest_prefix(). Typetokenizerand**kwargsinencode(), preferably with a tokenizerProtocol. Add docstrings for exportedprefix_cache_enabled()and publicencode().As per coding guidelines, “Annotate every function” and “Use docstrings rather than comments for externally usable interfaces.”
Also applies to: 67-70, 96-96
🤖 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/inputs/prefix_token_cache.py` at line 30, Complete the cache API contracts in _bucket_key(), _find_longest_prefix(), encode(), and prefix_cache_enabled(): add explicit return annotations, type encode()’s tokenizer and **kwargs (prefer a tokenizer Protocol), and document the exported prefix_cache_enabled() and public encode() APIs with docstrings.Source: Coding guidelines
tensorrt_llm/_torch/pyexecutor/py_executor.py (1)
3422-3427: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the unreachable disabled implementations.
Both methods return before their former logic. Delete the unreachable code and remove now-unused admission helpers if no other caller needs them.
tensorrt_llm/_torch/pyexecutor/py_executor.py#L3422-L3427: remove the unreachable admission-controller branch.tensorrt_llm/_torch/pyexecutor/py_executor.py#L3521-L3536: remove the unreachable idle-progress branch.🤖 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/pyexecutor/py_executor.py` around lines 3422 - 3427, Remove the unreachable admission-controller branch at tensorrt_llm/_torch/pyexecutor/py_executor.py:3422-3427 and the unreachable idle-progress branch at tensorrt_llm/_torch/pyexecutor/py_executor.py:3521-3536 from their respective methods, preserving the active return behavior. Afterward, remove any admission helpers left unused if no other callers depend on them.tensorrt_llm/_torch/disaggregation/native/bounce/config.py (1)
121-123: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse Python 3.10 annotation syntax consistently.
Replace
Optional[int]withint | None. ReplaceList[int]withlist[int].
tensorrt_llm/_torch/disaggregation/native/bounce/config.py#L121-L123: useint | Noneformin_blocksandmin_bytes.tensorrt_llm/_torch/disaggregation/native/bounce/impl.py#L67-L72: uselist[int]forblock_bytes_per_group.tensorrt_llm/_torch/disaggregation/native/bounce/impl.py#L102-L114: uselist[int]andfloat | Nonestyle where applicable.🤖 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/disaggregation/native/bounce/config.py` around lines 121 - 123, Use Python 3.10 union and built-in generic annotation syntax consistently: in config.py lines 121-123, update config_from_size parameters min_blocks and min_bytes to int | None; in impl.py lines 67-72, change block_bytes_per_group to list[int]; and in impl.py lines 102-114, replace applicable List[int] and Optional[float] annotations with list[int] and float | None.Sources: Coding guidelines, Learnings
🤖 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/inputs/prefix_token_cache.py`:
- Around line 42-53: Update the prefix-cache configuration flow around __init__
and the environment parsing at the referenced locations to safely parse invalid
numeric values with defaults before cache creation. Validate max_entries,
overlap, resync, min_chars, and bucket_chars constructor bounds before assigning
them, explicitly rejecting overlap <= 0, while preserving valid configured
values and preventing cache initialization from raising.
- Around line 96-151: Add direct PrefixTokenCache tests covering exact token
IDs, seam-resynchronization fallback, concurrent access, eviction, and invalid
settings. In PrefixTokenCache.encode, update the cache entry’s position in
_order whenever an existing entry is reused, moving its ID to the most-recent
end before _evict runs, so eviction follows LRU rather than FIFO.
In `@tensorrt_llm/inputs/registry.py`:
- Around line 141-147: Update the exception handling around
PrefixTokenCache.encode in the prefix-cache branch to catch only the dedicated
cache-fallback exception or explicitly enumerated recoverable exceptions.
Preserve the existing cache-disable fallback for those failures, while allowing
programming errors and unexpected tokenizer exceptions to propagate.
---
Nitpick comments:
In `@tensorrt_llm/_torch/disaggregation/native/bounce/config.py`:
- Around line 121-123: Use Python 3.10 union and built-in generic annotation
syntax consistently: in config.py lines 121-123, update config_from_size
parameters min_blocks and min_bytes to int | None; in impl.py lines 67-72,
change block_bytes_per_group to list[int]; and in impl.py lines 102-114, replace
applicable List[int] and Optional[float] annotations with list[int] and float |
None.
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 3422-3427: Remove the unreachable admission-controller branch at
tensorrt_llm/_torch/pyexecutor/py_executor.py:3422-3427 and the unreachable
idle-progress branch at tensorrt_llm/_torch/pyexecutor/py_executor.py:3521-3536
from their respective methods, preserving the active return behavior. Afterward,
remove any admission helpers left unused if no other callers depend on them.
In `@tensorrt_llm/inputs/prefix_token_cache.py`:
- Line 30: Complete the cache API contracts in _bucket_key(),
_find_longest_prefix(), encode(), and prefix_cache_enabled(): add explicit
return annotations, type encode()’s tokenizer and **kwargs (prefer a tokenizer
Protocol), and document the exported prefix_cache_enabled() and public encode()
APIs with docstrings.
🪄 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: 6c97242f-8dfe-4002-8997-496531318a02
📒 Files selected for processing (11)
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.pytensorrt_llm/_torch/attention_backend/sparse/dsa.pytensorrt_llm/_torch/disaggregation/native/bounce/config.pytensorrt_llm/_torch/disaggregation/native/bounce/impl.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/inputs/prefix_token_cache.pytensorrt_llm/inputs/registry.pytests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yamltests/integration/defs/disaggregated/test_disaggregated.pytests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.pytests/unittest/disaggregated/test_bounce.py
Comment on lines
+42
to
+53
| def __init__(self, | ||
| max_entries: int = 512, | ||
| overlap: int = 64, | ||
| resync: int = 32, | ||
| min_chars: int = 4096, | ||
| bucket_chars: int = 2048) -> None: | ||
| self._lock = threading.Lock() | ||
| self._max_entries = max_entries | ||
| self._overlap = overlap | ||
| self._resync = resync | ||
| self._min_chars = min_chars | ||
| self._bucket_chars = bucket_chars |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Validate prefix-cache settings before cache creation.
An invalid numeric environment value raises from int() during lazy cache creation. DefaultInputProcessor creates the cache before its cache fallback, so this fails the request when TLLM_PREFIX_TOKEN_CACHE=1.
Also reject overlap <= 0. With overlap == 0, i equals len(pstarts) for a cached prefix, and pstarts[i] raises IndexError.
Use safe parsing with defaults for environment values. Validate constructor bounds before storing the settings.
Also applies to: 102-104, 165-175
🤖 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/inputs/prefix_token_cache.py` around lines 42 - 53, Update the
prefix-cache configuration flow around __init__ and the environment parsing at
the referenced locations to safely parse invalid numeric values with defaults
before cache creation. Validate max_entries, overlap, resync, min_chars, and
bucket_chars constructor bounds before assigning them, explicitly rejecting
overlap <= 0, while preserving valid configured values and preventing cache
initialization from raising.
Comment on lines
+96
to
+151
| def encode(self, tokenizer, text: str, **kwargs) -> List[int]: | ||
| reuse, start_char, prev = 0, 0, None | ||
| with self._lock: | ||
| entry = self._find_longest_prefix(text) | ||
| if entry is not None: | ||
| ptext, pids, pends, pstarts, _ = entry | ||
| i = bisect.bisect_right(pends, len(ptext)) - self._overlap | ||
| if i > 0: | ||
| reuse, start_char = i, pstarts[i] | ||
| prev = (pids, pends, pstarts) | ||
|
|
||
| # tokenize outside the lock: this is the expensive part | ||
| enc = tokenizer(text[start_char:], | ||
| add_special_tokens=False, | ||
| return_offsets_mapping=True, | ||
| **kwargs) | ||
| new_ids, new_offsets = enc["input_ids"], enc["offset_mapping"] | ||
|
|
||
| if reuse: | ||
| pids, pends, pstarts = prev | ||
| span = min(self._resync, len(pids) - reuse, len(new_ids)) | ||
| if span <= 0 or list(new_ids[:span]) != list(pids[reuse:reuse + | ||
| span]): | ||
| # the tokenizer did not re-synchronize at the seam | ||
| with self._lock: | ||
| self.resync_failures += 1 | ||
| reuse, start_char = 0, 0 | ||
| enc = tokenizer(text, | ||
| add_special_tokens=False, | ||
| return_offsets_mapping=True, | ||
| **kwargs) | ||
| new_ids, new_offsets = enc["input_ids"], enc["offset_mapping"] | ||
|
|
||
| if reuse: | ||
| pids, pends, pstarts = prev | ||
| ids = pids[:reuse] + list(new_ids) | ||
| starts = pstarts[:reuse] + [a + start_char for a, _ in new_offsets] | ||
| ends = pends[:reuse] + [b + start_char for _, b in new_offsets] | ||
| else: | ||
| ids = list(new_ids) | ||
| starts = [a for a, _ in new_offsets] | ||
| ends = [b for _, b in new_offsets] | ||
|
|
||
| with self._lock: | ||
| if reuse: | ||
| self.hits += 1 | ||
| else: | ||
| self.misses += 1 | ||
| eid = self._next_id | ||
| self._next_id += 1 | ||
| key = self._bucket_key(text) | ||
| self._entries[eid] = (text, ids, ends, starts, key) | ||
| self._buckets.setdefault(key, []).append(eid) | ||
| self._order.append(eid) | ||
| self._evict() | ||
| return ids |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline tensorrt_llm/inputs/prefix_token_cache.py --items all
rg -n -C 3 --glob '*.py' \
'PrefixTokenCache|get_prefix_token_cache|TLLM_PREFIX_TOKEN_CACHE|resync_failures' \
tests tensorrt_llm || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 6684
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- candidate files ---'
git ls-files | rg '(^|/)(test|tests|testing)(/|_)|prefix_token_cache|CODING_GUIDELINES\.md$' | head -300
echo '--- implementation ---'
cat -n tensorrt_llm/inputs/prefix_token_cache.py | sed -n '1,192p'
echo '--- cache-related test symbols and strings ---'
rg -n -i --glob '*.py' \
'prefix.?token.?cache|resync.?fail|TLLM_PREFIX_TOKEN_CACHE|cache.*(evict|overlap|resync)|overlap|bucket_chars|max_entries|min_chars' \
. | head -500Repository: NVIDIA/TensorRT-LLM
Length of output: 17502
🏁 Script executed:
#!/bin/bash
set -u
echo '--- implementation ---'
cat -n tensorrt_llm/inputs/prefix_token_cache.py | sed -n '1,192p'
echo '--- exact cache-related references ---'
rg -n -i --glob '*.py' \
'prefix.?token.?cache|resync.?fail|TLLM_PREFIX_TOKEN_CACHE|cache.*(evict|overlap|resync)|bucket_chars|max_entries|min_chars' \
tests tensorrt_llm 2>/dev/null || true
echo '--- test files with likely cache names ---'
git ls-files | rg -i '(^|/)(test[^/]*)?prefix[^/]*cache|prefix[^/]*cache[^/]*test|cache[^/]*test.*prefix' || true
echo '--- Python test files mentioning tokenizer or input registry ---'
rg -l --glob 'test*.py' --glob '*.py' \
'input_processor|InputRegistry|tokenizer' tests 2>/dev/null | head -100 || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 31236
🏁 Script executed:
#!/bin/bash
set -u
echo '--- target test imports and body ---'
cat -n tests/unittest/disaggregated/test_coordinator_worker.py | sed -n '1,80p;360,455p'
echo '--- all test references to the exact function/class names ---'
rg -n -C 5 \
'test_prefix_token_cache_retokenizes_extended_text|PrefixTokenCache|prefix_cache_enabled|get_prefix_token_cache|resync_failures' \
tests --glob '*.py' 2>/dev/null || true
echo '--- cache implementation callers ---'
cat -n tensorrt_llm/inputs/registry.py | sed -n '110,165p'Repository: NVIDIA/TensorRT-LLM
Length of output: 12406
🏁 Script executed:
#!/bin/bash
set -u
echo '--- block-hash prefix-cache implementation ---'
rg -n -C 12 \
'def _encode_with_prefix_cache|_encode_with_prefix_cache|class BlockHashMixin' \
tensorrt_llm tests --glob '*.py' 2>/dev/null || true
echo '--- direct PrefixTokenCache imports/usages ---'
rg -n -C 8 \
'from .*prefix_token_cache|import .*prefix_token_cache|PrefixTokenCache\(' \
tests tensorrt_llm --glob '*.py' 2>/dev/null || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 11312
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
source_path = Path("tensorrt_llm/inputs/prefix_token_cache.py")
test_path = Path("tests/unittest/disaggregated/test_coordinator_worker.py")
source = ast.parse(source_path.read_text())
tests = ast.parse(test_path.read_text())
class_names = {
node.name
for node in ast.walk(tests)
if isinstance(node, ast.ClassDef)
}
direct_cache_calls = [
(node.lineno, ast.unparse(node))
for node in ast.walk(tests)
if isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == "PrefixTokenCache"
]
encode = next(
node for node in source.body
if isinstance(node, ast.ClassDef) and node.name == "PrefixTokenCache"
for node in node.body
if isinstance(node, ast.FunctionDef) and node.name == "encode"
)
evict = next(
node for node in ast.walk(source)
if isinstance(node, ast.FunctionDef) and node.name == "_evict"
)
order_writes_in_encode = [
(node.lineno, ast.unparse(node))
for node in ast.walk(encode)
if isinstance(node, (ast.Assign, ast.AnnAssign, ast.AugAssign))
and "_order" in ast.unparse(node)
]
evict_front_operations = [
(node.lineno, ast.unparse(node))
for node in ast.walk(evict)
if isinstance(node, ast.Call)
and isinstance(node.func, ast.Attribute)
and node.func.attr == "pop"
]
print(f"direct PrefixTokenCache test calls: {direct_cache_calls}")
print(f"test class names containing cache implementation: {sorted(name for name in class_names if 'Cache' in name)}")
print(f"encode _order writes: {order_writes_in_encode}")
print(f"_evict pop operations: {evict_front_operations}")
env_ints = [
(node.lineno, ast.unparse(node))
for node in ast.walk(source)
if isinstance(node, ast.Call)
and isinstance(node.func, NameError if False else ast.Name)
and node.func.id == "int"
]
print(f"integer environment conversions: {env_ints}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 729
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
source_path = Path("tensorrt_llm/inputs/prefix_token_cache.py")
test_path = Path("tests/unittest/disaggregated/test_coordinator_worker.py")
source = ast.parse(source_path.read_text())
tests = ast.parse(test_path.read_text())
direct_cache_calls = [
(node.lineno, ast.unparse(node))
for node in ast.walk(tests)
if isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == "PrefixTokenCache"
]
encode = next(
node for node in source.body
if isinstance(node, ast.ClassDef) and node.name == "PrefixTokenCache"
for node in node.body
if isinstance(node, ast.FunctionDef) and node.name == "encode"
)
evict = next(
node for node in ast.walk(source)
if isinstance(node, ast.FunctionDef) and node.name == "_evict"
)
order_writes_in_encode = [
(node.lineno, ast.unparse(node))
for node in ast.walk(encode)
if isinstance(node, (ast.Assign, ast.AnnAssign, ast.AugAssign))
and "_order" in ast.unparse(node)
]
evict_front_operations = [
(node.lineno, ast.unparse(node))
for node in ast.walk(evict)
if isinstance(node, ast.Call)
and isinstance(node.func, ast.Attribute)
and node.func.attr == "pop"
]
print(f"direct PrefixTokenCache test calls: {direct_cache_calls}")
print(f"encode _order writes: {order_writes_in_encode}")
print(f"_evict pop operations: {evict_front_operations}")
env_ints = [
(node.lineno, ast.unparse(node))
for node in ast.walk(source)
if isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == "int"
]
print(f"integer environment conversions: {env_ints}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 676
Add direct PrefixTokenCache tests and fix LRU recency.
The existing test covers BlockHashMixin’s separate cache. Add tests for exact IDs, seam fallback, concurrency, eviction, and invalid settings. _order.pop(0) plus no hit recency update implements FIFO, not LRU.
🤖 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/inputs/prefix_token_cache.py` around lines 96 - 151, Add direct
PrefixTokenCache tests covering exact token IDs, seam-resynchronization
fallback, concurrent access, eviction, and invalid settings. In
PrefixTokenCache.encode, update the cache entry’s position in _order whenever an
existing entry is reused, moving its ID to the most-recent end before _evict
runs, so eviction follows LRU rather than FIFO.
Comment on lines
+141
to
+147
| if use_cache: | ||
| with nvtx_range_debug("tokenize prompt (prefix cache)"): | ||
| try: | ||
| return cache.encode(self.tokenizer, prompt), None | ||
| except Exception: | ||
| # never fail a request over a cache problem | ||
| use_cache = False |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline tensorrt_llm/inputs/registry.py --items all
rg -n -C 4 --glob '*.py' \
'return_offsets_mapping|is_fast|PrefixTokenCache|prefix cache' \
tensorrt_llm tests || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 22809
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- registry.py ---'
sed -n '95,170p' tensorrt_llm/inputs/registry.py
printf '%s\n' '--- prefix_token_cache.py ---'
sed -n '1,190p' tensorrt_llm/inputs/prefix_token_cache.py
printf '%s\n' '--- cache callers and tests ---'
rg -n -C 5 --glob '*.py' \
'get_prefix_token_cache|PrefixTokenCache|resync_failures|return_offsets_mapping|TLLM_PREFIX_TOKEN_CACHE' \
tensorrt_llm testsRepository: NVIDIA/TensorRT-LLM
Length of output: 20163
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
for name in ("tensorrt_llm/inputs/registry.py",
"tensorrt_llm/inputs/prefix_token_cache.py"):
tree = ast.parse(Path(name).read_text(), filename=name)
print(f"--- {name}: exception handlers and raises ---")
for node in ast.walk(tree):
if isinstance(node, ast.ExceptHandler):
caught = ast.unparse(node.type) if node.type else "bare"
print(f"except {caught} at line {node.lineno}")
elif isinstance(node, ast.Raise):
print(f"raise at line {node.lineno}: {ast.unparse(node)}")
print(f"--- {name}: calls in PrefixTokenCache.encode ---")
if name.endswith("prefix_token_cache.py"):
for node in ast.walk(tree):
if isinstance(node, ast.FunctionDef) and node.name == "encode":
for call in ast.walk(node):
if isinstance(call, ast.Call):
print(f"line {call.lineno}: {ast.unparse(call)}")
PY
printf '%s\n' '--- focused diff ---'
git diff --unified=12 -- tensorrt_llm/inputs/registry.py tensorrt_llm/inputs/prefix_token_cache.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 4232
Narrow the prefix-cache exception handler.
PrefixTokenCache.encode() does not define a cache-specific fallback exception. except Exception can hide programming errors and unexpected tokenizer failures as cache misses. Catch a dedicated fallback exception or enumerate the concrete recoverable exceptions. Let other exceptions propagate.
🧰 Tools
🪛 Ruff (0.16.1)
[warning] 145-145: Do not catch blind exception: Exception
(BLE001)
🤖 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/inputs/registry.py` around lines 141 - 147, Update the exception
handling around PrefixTokenCache.encode in the prefix-cache branch to catch only
the dedicated cache-fallback exception or explicitly enumerated recoverable
exceptions. Preserve the existing cache-disable fallback for those failures,
while allowing programming errors and unexpected tokenizer exceptions to
propagate.
Sources: Coding guidelines, Linters/SAST tools
fredricz-20070104
left a comment
fredricz-20070104
left a comment
Collaborator
There was a problem hiding this comment.
Review summary - CONCERNS
Verdict: Cannot merge as-is — mergeable_state is dirty (rebase needed), and the PR bundles several unrelated, untested behavioural changes on top of the masked-indexer feature. No proven crash, but the risk-without-test surface is large.
Concerns
-
[MAJOR]
tensorrt_llm/inputs/prefix_token_cache.py- new 177-line cache has no tests- What is wrong:
encode()/_find_longest_prefix()/_evict()implement a thread-safe LRU + BPE-seam resync splice, but no test file for this module is in the diff. - How it fails: with
TLLM_PREFIX_TOKEN_CACHE=1, a splice/resync edge case (e.g.len(pids)-reuse <= 0, non-contiguous seam offsets) would silently return wrong token IDs for any prompt ≥min_chars, corrupting generation. Off-by-default limits blast radius but the path is unexercised. - Suggested fix: add unit tests (hit/miss, resync fallback, eviction, byte-identical vs. full tokenization). This module is also unrelated to the PR title — consider splitting it out.
- What is wrong:
-
[MAJOR]
tensorrt_llm/_torch/pyexecutor/py_executor.py:3422&:3521- disagg admission/idle collectives disabled, untested- What is wrong:
_apply_disagg_transfer_admissionnow unconditionally returns(requests, False)and_check_disagg_transfer_progress_when_idlereturns immediately, removing the in-flight-transfer bound and two rank-scoped collectives. - How it fails: correctness (rank-uniformity, completion reaped elsewhere) is asserted only in comments. If any rank ever still needed the status wait, dropping the WORLD/TP collective can hang/diverge disagg ranks; unbounded admission can pressure transfer buffers under load. No test covers this, and it is unrelated to the feature.
- Suggested fix: cover with a disagg integration test or gate behind a flag; at minimum delete the now-unreachable code below each
return.
- What is wrong:
-
[MAJOR]
tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py:326- masked transfer only tested at page-table level- What is wrong: the masked REPLICATED indexer view is validated by construction only; the PR's claim that the masked subset "transfers correctly, including under PP resharding" has no end-to-end test.
- How it fails: a per-pool matching mismatch for the masked subset under PP resharding would misroute/drop indexer K-cache and go uncaught.
- Suggested fix: add a masked-layout analogue of
test_v1_dsa_indexer_replicated_transfer_across_pp.
Minor notes (non-blocking)
py_executor.py:3428- remove the dead code after the early returns.config.py:118- default gate semantics change (min_blocks 96→1, new min_bytes=2 MiB); flag in release notes.kv_extractor.py:336- replacegetattron always-set manager attrs with direct access.test_extractor.py- annotateindexer_k_cache_layer_maskaslist[bool] | None.
QA view
- Test coverage: partial - dsa.py MTP-draft rebuild and bounce byte-gate are well covered;
prefix_token_cache.pyand the py_executor disabling are uncovered; masked indexer is construction-only. - SM coverage: DSA/GLM-5.2 masked path targets Hopper (sm90) and Blackwell (sm100/fp8); new DSA tests run
skip_pre_hopper+ DeepGEMM (Hopper only), no explicit Blackwell run for the masked path. - Test code: masked page-table test asserts shape only; no prefix-cache test; helper param unannotated.
- Test time: small - two Hopper-gated unit tests + CPU bounce/extractor cases; integration test only renamed an env var.
- Needs
/qa-verify: yes - disagg behavioural changes + preference flip (GLM 5.2 → PYTHON) + untested new module warrant a real GLM-5.2 disagg run (incl. PP resharding) before trusting.
Possible new issues
- Unbounded disagg admission may exhaust transfer buffers under concurrency.
- Removed status-vote collectives depend on rank-uniformity at every site; if violated, cross-rank hang.
- Byte-gate default change flips which transfers take the coalesced-bounce path for existing models.
What I could not verify
- Runtime safety of removing the disagg collectives (depends on invariants at other executor call sites not shown).
- Whether
get_indexer_k_cache_pool_layer_idx/local_layer_idsindexing holds under all PP configs. - End-to-end masked-indexer transfer correctness — only page-table construction is visible in the diff.
Automated review by NVCortex Lite, run by @fredricz-20070104.
fredricz-20070104
approved these changes
Aug 7, 2026
fredricz-20070104
left a comment
fredricz-20070104
left a comment
Collaborator
There was a problem hiding this comment.
Review summary - Approve (non-blocking)
Approving so this is not blocked on me. The points raised in my review comment above are non-blocking — please read them and address what you agree with before merging.
Worth doing before this is relied on: Bundles disagg-serving behavioural changes (admission control + collective removal) and a masked-indexer transceiver path with no end-to-end disagg test, plus a preference flip (GLM 5.2 now PYTHON) and an untested new tokenization cache; a human QA should run the GLM-5.2 disagg path (incl. PP resharding) before trusting this.
Automated review by NVCortex Lite, run by @fredricz-20070104.
yuxianq
reviewed
Aug 7, 2026
| # would resolve to req_idx == num_seqs and index out of bounds below. | ||
| # Only the device buffer is refreshed: host_req_idx_per_token is read | ||
| # solely by its own producer in prepare_for_indices_conversion(). | ||
| cu_seq_lens = torch.cumsum(seq_lens, dim=0, dtype=torch.int32) |
Collaborator
There was a problem hiding this comment.
#16925 has already fixed this issue, so we can rebase and revert changes in dsa.py and deepseek_v4.py.
yuxianq
reviewed
Aug 7, 2026
Collaborator
There was a problem hiding this comment.
If we don't need to update dsa.py, we also don't need to update test_dsa_indexer.py
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
The Python (v2) KV-cache transceiver previously raised
NotImplementedErrorfor the per-layer masked DSA indexer k-cache pool (cross-layer indexer sharing, e.g. GLM 5.2), so those checkpoints were forced onto the C++ transceiver — #16558 added aGlmMoeDsaForCausalLM -> CPPpreference as a stop-gap while the Python path lacked this support.This PR teaches the Python transceiver to handle the masked layout, mirroring the C++ support added in #16558:
build_page_table(tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py): the indexerREPLICATEDpool view now covers only the indexer-owning layers — onebuffer_entriesrow per owning layer, each mapped to its packed pool row viaimpl.get_indexer_k_cache_pool_layer_idx(lid)— and skips the indexer pool entirely for a layer group with no owning layers (which would otherwise hit the null-pool getter). The dense/unmasked layout is byte-for-byte unchanged.pool_role+global_layer_idoverlap, so a masked subset transfers correctly, including under PP resharding (the generic analogue of the C++indexerLayerNumPerPP/targetIRanksForIndexerKCacheinterval logic).GlmMoeDsaForCausalLMprefers the Python transceiver again like the other DeepSeek-family checkpoints, reverting theCPPoverride that [None][perf] Allocate DSA indexer k-cache only for layers that own an indexer #16558 added only because Python lacked masked-pool support.Test Coverage
tests/unittest/disaggregated/test_extractor.py::test_v1_dsa_masked_indexer_page_table_covers_owning_layers(new) — builds a V1KVCacheManagerwith a per-layer indexer mask ([True, False, True, False]) and asserts the indexer view isREPLICATED, covers exactly the two owning layers, and maps them to the correct packed pool rows/offsets.test_extractor.py::test_v1_dsa_indexer_page_table_is_replicated_with_per_layer_entriesand::test_v1_dsa_indexer_replicated_transfer_across_pp(existing) — continue to guard the dense layout and the end-to-end replicated transfer path.tests/unittest/llmapi/test_llm_args.py::TestDeepseekTransceiverPreference::test_preference_per_architecture— updated to expectPYTHONforglm_moe_dsaagain.PR Checklist
get_preferred_transceiver_runtimeis an internal preference hook).GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.Dev Engineer Review
build_page_tablenow supports masked DSA indexer K-cache pools.REPLICATEDviews only for indexer-owning layers.GlmMoeDsaForCausalLMnow prefers the Python transceiver.CODING_GUIDELINES.md.QA Engineer Review
GlmMoeDsaForCausalLM.test_on_update_kv_lens_rebuilds_req_idx_in_draft_looptest_on_update_kv_lens_matches_prepare_on_target_forwardtests/integration/test_lists/,test-db/, orqa/coverage was identified in the provided changes.