[None][chore] ReuseScope refactor (cp #14140) by lancelly · Pull Request #14353 · NVIDIA/TensorRT-LLM

lancelly · 2026-05-20T09:07:00Z

Summary

Stacked on top of #14351. Cherry-picks #14140 onto feat/deepseek_v4 and fixes two regressions that surface once that refactor lands on the DSv4 branch.

[None][chore] Refactor salting support for KVCacheManagerV2 #14140 (cherry-picked) — Refactor cache_salt_id / lora_task_id into a ReuseScope NamedTuple. Eliminates positional (int | None, int | None) parameters across the radix tree API. Conflict resolution against the DSv4 branch:
- _kv_cache_manager.py: collapse the public create_kv_cache / private _create_kv_cache split into a single entry point ([None][chore] Refactor salting support for KVCacheManagerV2 #14140's shape), but keep the DSv4-only expected_prompt_length parameter.
- _kv_cache.py: drop cache_salt_id from __init__ (now folded into reuse_scope), keep expected_prompt_length.
- resource_manager.py: switch caller from impl._create_kv_cache(lora_task_id, ..., cache_salt_id=, expected_prompt_length=) to impl.create_kv_cache(ReuseScope(lora_id=..., salt=...), input_tokens, id=request_id, expected_prompt_length=...). Preserves the DSv4 IndexMapper free-slot guard.
Re-migrate router callers from TreeTaskId → ReuseScope. PR [None][chore] Replace cache_salt_id impl with main's TreeTaskId (cp #13793) #14351's eaddc9ac22 migrated callers to a TreeTaskId seed (_KVCache._make_tree_task_id(None, cache_salt_id)); [None][chore] Refactor salting support for KVCacheManagerV2 #14140 supersedes that API, so two sites in tensorrt_llm/serve/router.py and four sites in tests/unittest/disaggregated/test_router.py need to be re-migrated to ReuseScope(salt=cache_salt_id) (or ReuseScope() for the unsalted path). Drops the now-unused _KVCache import.
Event-hash V1 compat fix (addresses @peihu-nv's review feedback on [None][chore] Replace cache_salt_id impl with main's TreeTaskId (cp #13793) #14351). After [None][chore] Refactor salting support for KVCacheManagerV2 #14140, RootBlock no longer exposes lora_task_id / cache_salt_id directly — both fields live on root.reuse_scope. KVCacheEventManager._root_attrs_from_root_block was still reading them via getattr(root, "lora_task_id", None) / getattr(root, "cache_salt_id", None), so every emitted V1-compat event silently received (None, None) and the hash collapsed to the same value for any LoRA/salt request. Downstream Dynamo prefix-cache routing depends on those hashes, so this regression would materially degrade hit rate. Fix reads root.reuse_scope first and falls back to the legacy attributes (keeps any in-flight pre-refactor RootBlock instances working). Includes a regression test (test_v2_kv_cache_event_manager_v1_hash_reads_root_reuse_scope) covering: (a) ReuseScope-shaped vs legacy root produce identical V1 hashes for the same lora/salt, (b) different scopes still produce different hashes (no silent collapse), (c) empty ReuseScope matches a legacy root with no salt/lora.

After NVIDIA#14140 (cherry-picked in NVIDIA#14353) RootBlock no longer exposes ``lora_task_id`` / ``cache_salt_id`` directly — both fields are folded into a ``ReuseScope`` NamedTuple attached as ``root.reuse_scope``. ``KVCacheEventManager._root_attrs_from_root_block`` was still reading the legacy attribute names via ``getattr(..., None)``, so every emitted event silently received ``(None, None)`` and the V1-compat hash collapsed to the same value for any LoRA/salt request. Downstream Dynamo routing depends on those hashes, so this regression would materially degrade prefix cache hit rate for any request carrying a LoRA task id or a cache salt. Fix ``_root_attrs_from_root_block`` to prefer ``root.reuse_scope`` and fall back to the legacy attributes (keeps any in-flight pre-refactor RootBlock instances working). Also: - Update the now-stale ``test_v2_root_key_distinguishes_lora_from_cache_salt_id`` to use the new single-arg ``RootBlock.make_key(ReuseScope(...))`` API. - Extend ``_FakeRootBlock`` with an optional ``reuse_scope`` kwarg so tests can mimic both the pre- and post-refactor RootBlock shape. - Add a regression test ``test_v2_kv_cache_event_manager_v1_hash_reads_root_reuse_scope`` that asserts (a) a ReuseScope-shaped root and a legacy root with the same lora/salt produce identical V1 event hashes, (b) different scopes still produce different hashes (no silent collapse), and (c) an empty ReuseScope matches a legacy root with no salt/lora. Addresses reviewer feedback on NVIDIA#14351 (peihu-nv). Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>

lancelly · 2026-05-21T02:52:40Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-21T02:59:01Z

PR_Github #49555 [ run ] triggered by Bot. Commit: 8003b8a Link to invocation

tensorrt-cicd · 2026-05-21T06:55:36Z

PR_Github #49555 [ run ] completed with state SUCCESS. Commit: 8003b8a
/LLM/main/L0_MergeRequest_PR pipeline #39183 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

lancelly · 2026-05-21T07:58:41Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-21T08:05:34Z

PR_Github #49644 [ run ] triggered by Bot. Commit: 63732b3 Link to invocation

tensorrt-cicd · 2026-05-21T12:29:19Z

PR_Github #49644 [ run ] completed with state SUCCESS. Commit: 63732b3
/LLM/main/L0_MergeRequest_PR pipeline #39262 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

lancelly · 2026-05-21T13:19:24Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-21T13:26:11Z

PR_Github #49705 [ run ] triggered by Bot. Commit: 63732b3 Link to invocation

tensorrt-cicd · 2026-05-21T14:42:17Z

PR_Github #49705 [ run ] completed with state SUCCESS. Commit: 63732b3
/LLM/main/L0_MergeRequest_PR pipeline #39311 completed with status: 'SUCCESS'

CI Report

Link to invocation

…4140) Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

After the ReuseScope refactor cherry-picked into this stack, ``RootBlock.make_key`` takes a single ``ReuseScope`` instead of the legacy ``(lora_task_id, cache_salt_id)`` pair, and ``sequence_to_blockchain_keys`` likewise drops the trailing ``cache_salt_id`` argument in favour of a ``ReuseScope`` positional. Two callers on the base feat/deepseek_v4 branch still used the legacy shapes and would trip the seven disaggregated/test_router.py tests that exercise the V2 hash path: - tensorrt_llm/serve/router.py — two call sites passed ``(None, cache_salt_id)`` to V2RootBlock.make_key - tests/unittest/disaggregated/test_router.py — four call sites passed raw ``None`` or a trailing ``cache_salt_id`` to sequence_to_blockchain_keys Migrate every call to ``ReuseScope(salt=cache_salt_id)`` (or ``ReuseScope()`` for the unsalted path) so the V2 prefix-hash flow matches what the radix tree now expects. Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>

After NVIDIA#14140 (cherry-picked in NVIDIA#14353) RootBlock no longer exposes ``lora_task_id`` / ``cache_salt_id`` directly — both fields are folded into a ``ReuseScope`` NamedTuple attached as ``root.reuse_scope``. ``KVCacheEventManager._root_attrs_from_root_block`` was still reading the legacy attribute names via ``getattr(..., None)``, so every emitted event silently received ``(None, None)`` and the V1-compat hash collapsed to the same value for any LoRA/salt request. Downstream Dynamo routing depends on those hashes, so this regression would materially degrade prefix cache hit rate for any request carrying a LoRA task id or a cache salt. Fix ``_root_attrs_from_root_block`` to prefer ``root.reuse_scope`` and fall back to the legacy attributes (keeps any in-flight pre-refactor RootBlock instances working). Also: - Update the now-stale ``test_v2_root_key_distinguishes_lora_from_cache_salt_id`` to use the new single-arg ``RootBlock.make_key(ReuseScope(...))`` API. - Extend ``_FakeRootBlock`` with an optional ``reuse_scope`` kwarg so tests can mimic both the pre- and post-refactor RootBlock shape. - Add a regression test ``test_v2_kv_cache_event_manager_v1_hash_reads_root_reuse_scope`` that asserts (a) a ReuseScope-shaped root and a legacy root with the same lora/salt produce identical V1 event hashes, (b) different scopes still produce different hashes (no silent collapse), and (c) an empty ReuseScope matches a legacy root with no salt/lora. Addresses reviewer feedback on NVIDIA#14351 (peihu-nv). Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>

lancelly · 2026-05-21T14:52:56Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-21T15:00:21Z

PR_Github #49720 [ run ] triggered by Bot. Commit: 4bacb3c Link to invocation

peihu-nv

LGTM.

Minor follow-up later: a couple tests not in CI still use stale APIs (cache_salt_id= / make_key(None)).

tensorrt-cicd · 2026-05-21T18:11:01Z

PR_Github #49720 [ run ] completed with state SUCCESS. Commit: 4bacb3c
/LLM/main/L0_MergeRequest_PR pipeline #39326 completed with status: 'SUCCESS'

CI Report

Link to invocation

Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com> Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com> Co-authored-by: Yao Yao <lowsfer@users.noreply.github.com>

github-actions Bot assigned lancelly May 20, 2026

lancelly changed the title ~~[None][chore] DSv4: ReuseScope refactor + KV cache reuse probe (cp #14140 + #14333)~~ [None][chore] ReuseScope refactor + KV cache reuse probe (cp #14140 + #14333) May 21, 2026

lancelly force-pushed the feat/dsv4-refactor-and-probe branch from c5bc871 to a4c6725 Compare May 21, 2026 01:41

lancelly force-pushed the feat/dsv4-refactor-and-probe branch from a4c6725 to 29401a9 Compare May 21, 2026 01:54

lancelly force-pushed the feat/dsv4-refactor-and-probe branch from 4c1b768 to dc137ab Compare May 21, 2026 02:38

lancelly changed the title ~~[None][chore] ReuseScope refactor + KV cache reuse probe (cp #14140 + #14333)~~ [None][chore] ReuseScope refactor (cp #14140) + DSv4 follow-ups May 21, 2026

lancelly force-pushed the feat/dsv4-refactor-and-probe branch from dc137ab to 8003b8a Compare May 21, 2026 02:48

lancelly marked this pull request as ready for review May 21, 2026 02:52

lancelly requested review from a team as code owners May 21, 2026 02:52

lancelly requested review from QiJune, kaiyux and syuoni and removed request for a team May 21, 2026 02:52

lancelly mentioned this pull request May 21, 2026

[None][feat] KV reuse probe #14395

Merged

lfr-0531 added the deepseek-v4 label May 21, 2026

lancelly changed the title ~~[None][chore] ReuseScope refactor (cp #14140) + DSv4 follow-ups~~ [None][chore] ReuseScope refactor (cp #14140) May 21, 2026

lowsfer and others added 3 commits May 21, 2026 07:45

[None][chore] Refactor salting support for KVCacheManagerV2 (NVIDIA#1…

346770f

…4140) Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

lancelly mentioned this pull request May 21, 2026

[None][fix] adp_router: salt-aware probe_prefix_match_length on v1/v2 KV cache backends #14417

Closed

2 tasks

lancelly force-pushed the feat/dsv4-refactor-and-probe branch from 63732b3 to 4bacb3c Compare May 21, 2026 14:49

peihu-nv approved these changes May 21, 2026

View reviewed changes

pcastonguay approved these changes May 21, 2026

View reviewed changes

peihu-nv merged commit a2842b1 into NVIDIA:feat/deepseek_v4 May 21, 2026
6 checks passed

This was referenced May 22, 2026

[None][feat] adp_router: salt-aware probe_prefix_match_length on v1/v2 KV cache backends #14435

Merged

[None][fix] Migrate stale event-manager test calls to ReuseScope API #14445

Merged

Uh oh!

Conversation

lancelly commented May 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Uh oh!

lancelly commented May 21, 2026

Uh oh!

tensorrt-cicd commented May 21, 2026

Uh oh!

tensorrt-cicd commented May 21, 2026

Uh oh!

lancelly commented May 21, 2026

Uh oh!

tensorrt-cicd commented May 21, 2026

Uh oh!

tensorrt-cicd commented May 21, 2026

Uh oh!

lancelly commented May 21, 2026

Uh oh!

tensorrt-cicd commented May 21, 2026

Uh oh!

tensorrt-cicd commented May 21, 2026

Uh oh!

lancelly commented May 21, 2026

Uh oh!

tensorrt-cicd commented May 21, 2026

Uh oh!

peihu-nv left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented May 21, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

lancelly commented May 20, 2026 •

edited

Loading