From 8c8570801dd46bf1d683a3991f24798c68f9e1c4 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 21 Jul 2026 13:10:19 +0800 Subject: [PATCH 01/22] [TRTLLM-11875][feat] Add V2 Mamba snapshot reuse (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- PR3_MAMBA_V2_REVIEW.md | 239 +++ docs/source/developer-guide/telemetry.md | 4 +- docs/source/features/kvcache.md | 34 + tensorrt_llm/_torch/pyexecutor/_util.py | 213 ++- .../_torch/pyexecutor/kv_cache_manager_v2.py | 19 +- .../_torch/pyexecutor/mamba_cache_manager.py | 1668 ++++++++++++++++- .../_torch/pyexecutor/py_executor_creator.py | 5 +- .../_torch/pyexecutor/resource_manager.py | 26 +- tensorrt_llm/llmapi/__init__.py | 7 +- tensorrt_llm/llmapi/llm_args.py | 95 +- .../runtime/kv_cache_manager_v2/__init__.pyi | 1 + .../kv_cache_manager_v2/_block_radix_tree.py | 5 +- .../runtime/kv_cache_manager_v2/_config.py | 2 + .../kv_cache_manager_v2/_storage_manager.py | 19 +- .../usage/llm_args_golden_manifest.json | 2 +- .../defs/accuracy/test_llm_api_pytorch.py | 26 +- .../executor/test_kv_cache_manager_v2.py | 18 + .../executor/test_kv_cache_v2_scheduler.py | 15 + .../executor/test_mamba_cache_manager.py | 1433 +++++++++++++- ...y_executor_creator_mla_cache_reuse_sync.py | 33 +- .../modeling/test_modeling_multimodal.py | 2 +- .../test_kv_cache_manager_v2.py | 228 +++ .../test_kv_cache_stats_behavior.py | 1 + tests/unittest/llmapi/test_llm_args.py | 118 +- 24 files changed, 4019 insertions(+), 194 deletions(-) create mode 100644 PR3_MAMBA_V2_REVIEW.md diff --git a/PR3_MAMBA_V2_REVIEW.md b/PR3_MAMBA_V2_REVIEW.md new file mode 100644 index 000000000000..e24ddad92013 --- /dev/null +++ b/PR3_MAMBA_V2_REVIEW.md @@ -0,0 +1,239 @@ + + +# PR3 V2 Mamba snapshot reuse review + +This document records the design decisions, code changes, experiments, and +verification for PR3. It is intended to make the final PR diff reviewable +against upstream `main`, which already contains PR1, rather than against an +older stacked-PR base. + +## Baseline + +- PR1 was squash-merged into upstream `main` as `bb90835f8a`. +- PR3 was rebuilt on the latest upstream `main` from its four PR3-specific + commits. +- The original PR1 commit chain and the two stacked-branch synchronization + merges were removed from PR3 history. + +The final PR3 diff must be reviewed relative to upstream `main`. + +## Review requirements + +1. Select the V1 Mamba manager by default. Select V2 only when + `use_kv_cache_manager_v2: true` is explicitly configured. +2. Move Mamba snapshot policy under `mamba_state_config`: + - `periodic_snapshot_interval` controls regular snapshot boundaries. + - `additional_snapshot_offsets_from_start` expresses fixed boundaries + measured from the prompt start. + - `additional_snapshot_offsets_from_end` expresses fixed boundaries + measured backward from the prompt end; zero selects the prompt end. +3. Reject unsupported combinations before manager construction: + - V1 supports periodic snapshots only. + - V2 does not support disaggregated serving. +4. Keep generic KV cache manager changes minimal and express Mamba behavior + through a specialized extension wherever possible. +5. Determine whether `commit_min_snapshot` removes the need for + `allow_prefix_sibling`, using multi-turn and forked-prefix reuse tests. +6. Explain and justify the extra attention-slot bound. + +## Configuration and manager-selection decisions + +`forward` and `backward` were rejected as field names because they can be +misread as execution or traversal directions. The selected names state both +the unit and reference anchor explicitly. + +| Serving mode | `use_kv_cache_manager_v2` | Additional offsets | Result | +| --- | --- | --- | --- | +| Aggregated | false or auto | Empty | V1 C++ manager | +| Aggregated | true | Any valid policy | V2 Mamba manager | +| Aggregated | false or auto | Non-empty | Configuration error | +| Disaggregated | false or auto | Empty | V1 C++ or mixed manager, according to the transceiver | +| Disaggregated | true | Any | Configuration error | +| Disaggregated | false or auto | Non-empty | Configuration error | + +The `TLLM_MAMBA_MANAGER_PREFERENCE=V2` override conflicts with the requirement +that V2 be selected only by the explicit public setting. The V2 override has +therefore been removed. An explicit V2 setting that conflicts with the CPP, +MIXED, or `TRTLLM_USE_PY_MAMBA` compatibility overrides is rejected rather +than silently routed to another implementation. + +The snapshot resolver uses these exact semantics: + +- `periodic_snapshot_interval=0` disables periodic snapshots; +- start offsets are strictly positive absolute token boundaries; +- end offsets are non-negative and resolve to `prompt_len - offset`, so zero + selects the final prompt boundary; +- resolved positions outside `(0, prompt_len]` are ignored; +- periodic and fixed boundaries are deduplicated and sorted. + +## Public configuration API + +This PR deliberately makes the snapshot-policy rename a breaking Python API +change. The former top-level `mamba_state_cache_interval` constructor argument +and Python property are removed without a model validator or deprecated alias, +and strict model construction rejects the removed field. YAML/JSON file loaders +retain a narrow compatibility migration to +`mamba_state_config.periodic_snapshot_interval`; setting both paths is an +error. All in-tree callers, examples, telemetry metadata, and canonical +documentation use the nested field. + +## KV cache extension audit + +The completed read-only audit found that the generic manager should retain only +two neutral PR3 changes: + +- a neutral constructor input for reserved index slots, used only to size the + stable `IndexMapper` storage; +- reporting every physical pool group in iteration statistics, which is a + generic correctness fix and needs a non-Mamba test. + +The remaining changes are Mamba-specific and have moved into +`V2MambaHybridCacheManager`, following `DeepseekV4CacheManager`: + +- CUDA-graph and attention-DP dummy-slot calculation; +- SSM/conv roles and page-table layout; +- SWA scratch layout for recurrent pools; +- snapshot commit, history, final-free, and context-update lifecycle; +- recurrent-state invalid-value checking. + +`KVCacheDesc.capacity` remains token capacity. The only Mamba-specific sizing +field is `num_ssm_slots`, which counts the live, snapshot, and dummy recurrent +state slots associated with that logical request. The storage planner sums it +for SSM lifecycles. If the batch has any non-live SSM capacity, the planner also +reserves one retained partial attention page per request lineage. There is no +separate `num_extra_attention_slots` field or layer-specific `BatchDesc` +descriptor. + +The subclass computes the exact CUDA-graph and attention-DP dummy count before +calling the base constructor, then passes that count through the neutral +`num_reserved_index_slots` extension. Its CUDA and pinned-host state-index +tensors are allocated once for `max_batch_size + reserved slots`; later dummy +insertion therefore cannot replace an aliased tensor or invalidate a captured +CUDA graph. The subclass also owns the mixed-pool page table, event-window +mapping for `SsmLayerConfig`, snapshot commit/history lifecycle, and state +diagnostics. The generic manager no longer knows any Mamba roles or policies. + +Metadata preparation still runs on attention-only pipeline-parallel ranks even +though no local Mamba kernel consumes state indices there. Both V1 and V2 now +return harmless zero placeholders on those ranks instead of consulting Mamba +state maps or buffers that are intentionally absent. The V2 invalid-value +diagnostic also scans every attention layer: a runtime layer group represents +a lifecycle, not necessarily one physical pool, so layers with different +buffer sizes cannot safely be deduplicated by group ID. + +The audit classified each generic KVCacheManagerV2 hunk as one of: + +- generic prerequisite already supplied by `commit_min_snapshot`; +- generic extension point required by Mamba; +- Mamba-only policy that belongs in `mamba_cache_manager.py`; +- obsolete code removable from PR3. + +## `allow_prefix_sibling` experiment + +The PR3-specific constructor flag, sibling-retention rule, alternate SSM +matcher, and both call sites were removed. `commit_min_snapshot` remains the +only runtime prerequisite. + +The first experiment exposed a generic replacement-order bug: when a longer +partial block replaced the only child of a root, `detach_next()` could prune +the now-empty root before the replacement was attached. The replacement is now +registered before covered children are detached. + +The resulting semantics are deliberately simpler: + +- monotonically extending snapshots at 10, 20, and 25 tokens reuse 10, 20, + and 25 tokens respectively; +- a fork sharing only 15 tokens cannot consume the retained 20-token state and + safely re-prefills from zero; +- a fork sharing 25 tokens reuses the 20-token state; +- a block-aligned 32-token snapshot remains reusable by a later fork. + +Only the latest partial snapshot within one token block and request lineage is +retained. An earlier same-block fork may therefore miss a reusable state, but +it cannot read a state computed from tokens beyond the fork boundary. This is +an accepted efficiency tradeoff for removing the custom sibling structure. + +## Extra attention slots + +These are physical attention-cache slots used +when a non-block-aligned Mamba snapshot preserves a copied partial attention +page in the radix tree. They are not request slots or CUDA-graph/attention-DP +dummy slots. Their capacity is additional to the normal attention-page budget +because the active request still owns a writable partial page while the radix +tree retains the snapshot copy. + +Each logical request descriptor budgets at least one live SSM state. The +planner uses `total_num_ssm_slots > num_requests` only to detect that the batch +has non-live SSM capacity. It then reserves one retained partial attention page +for every request lineage and adds that bound to every attention lifecycle; +nothing is supplied by an external caller. Covered partial siblings are removed +as a lineage advances, so more than one retained page per lineage is +unnecessary. + +This intentionally trades some capacity precision for a smaller interface. +The bound does not require `total_num_ssm_slots >= 2 * num_requests`. When the +number of non-live slots is smaller than the number of lineages, or when those +slots are aligned snapshots or reserved dummy states, it simply over-reserves +attention capacity. It never depends on snapshot alignment or slot +distribution. The normal request token capacity already accounts for a resumed +request's writable page; this bound covers the retained radix-tree snapshot +page. The V1 estimator retains its existing policy because this single-field +bound is specific to V2. + +## PP/spec layout and affine sizing + +One-model MTP with pipeline parallelism is expected to provide an explicit +`pp_partition` covering the base-model layers. The existing PP helper uses that +partition as the base/spec boundary and assigns appended speculative layers to +the last rank. Supporting an automatically derived boundary when +`pp_partition` is absent is intentionally left to a separate change. + +The hybrid affine estimator uses the same normalized masks and explicit PP +layout to count local attention and Mamba layers. This avoids two failures: + +- custom `pp_partition` being applied to an attention-only layer count and + raising during startup; +- a separate, attention-only MTP draft cache being charged for phantom target + Mamba state. + +Target-only and draft-only masks are shared between estimation and runtime +construction, so the budget split uses the physical layout each manager will +actually allocate. + +## Verification log + +- `python3 -m py_compile` passed after the configuration, routing, subclass, + and test migrations. +- Focused configuration/routing/snapshot-point tests: 34 passed. +- Full `test_mamba_cache_manager.py`: 85 passed. +- Three focused final-audit regressions passed: V1 and V2 metadata preparation + on ranks with no local Mamba layers, and invalid-value detection in a second + differently sized attention pool that shares a lifecycle. +- Six focused PP/MTP sizing cases passed across default/custom PP partitions, + Cpp/V2 estimators, target-only masks, and attention-only draft masks. +- The combined Mamba, cache-budget split, dual-pool, and executor-creator suite + passed: 151 tests. +- Generic KVCMv2 and executor-creator tests: 24 passed. +- Full `test_llm_args.py`: 255 passed, 3 skipped. A preceding attempt failed + because the environment's MPI session-reuse worker did not report an identity + within 60 seconds on either spawn attempt (`0/1` workers); it did not reach + Mamba configuration or manager code. +- The C++ estimator regression for disabled periodic snapshots on a + recurrent-only rank passed, as did the V2 partial-attention sizing test. +- Runtime radix-tree focused tests: 2 passed. +- Runtime `TestSSMSupport`: 13 passed. +- Runtime `TestNoBatching`: 24 passed, 12 skipped. +- `scripts/generate_llm_args_golden_manifest.py --check` passed. The manifest + exposes only `kv_cache_config.mamba_state_config.periodic_snapshot_interval`. +- The telemetry generator cannot load legacy `TrtLlmArgs` in this local + environment and would incorrectly erase that entire table. The generated + noise was discarded; the two existing Mamba telemetry rows were updated to + the manifest's canonical path without changing unrelated rows. +- Isort/YAPF formatting and the completed runtime experiment's lint/diff + checks passed. The final read-only audit found no remaining correctness + blockers after the PP/spec sizing fixes. +- Final manifest, syntax, diff, and full pre-commit checks passed before + publication. diff --git a/docs/source/developer-guide/telemetry.md b/docs/source/developer-guide/telemetry.md index fc0e1fc8bbe1..96b52ff4704f 100644 --- a/docs/source/developer-guide/telemetry.md +++ b/docs/source/developer-guide/telemetry.md @@ -122,7 +122,7 @@ unset or when the safety sanitizer rejects the runtime value. | `kv_cache_config.mamba_ssm_cache_dtype` | `Literal['auto', 'float16', 'bfloat16', 'float32']` | `categorical` | | `auto`, `float16`, `bfloat16`, `float32` | | `kv_cache_config.mamba_ssm_philox_rounds` | `` | `value` | | | | `kv_cache_config.mamba_ssm_stochastic_rounding` | `` | `value` | | | -| `kv_cache_config.mamba_state_cache_interval` | `` | `value` | | | +| `kv_cache_config.mamba_state_config.periodic_snapshot_interval` | `` | `value` | | | | `kv_cache_config.max_attention_window` | `Optional[List[int]]` | `value` | | | | `kv_cache_config.max_gpu_total_bytes` | `` | `value` | | | | `kv_cache_config.max_tokens` | `Optional[int]` | `value` | | | @@ -431,7 +431,7 @@ unset or when the safety sanitizer rejects the runtime value. | `kv_cache_config.mamba_ssm_cache_dtype` | `Literal['auto', 'float16', 'bfloat16', 'float32']` | `categorical` | | `auto`, `float16`, `bfloat16`, `float32` | | `kv_cache_config.mamba_ssm_philox_rounds` | `` | `value` | | | | `kv_cache_config.mamba_ssm_stochastic_rounding` | `` | `value` | | | -| `kv_cache_config.mamba_state_cache_interval` | `` | `value` | | | +| `kv_cache_config.mamba_state_config.periodic_snapshot_interval` | `` | `value` | | | | `kv_cache_config.max_attention_window` | `Optional[List[int]]` | `value` | | | | `kv_cache_config.max_gpu_total_bytes` | `` | `value` | | | | `kv_cache_config.max_tokens` | `Optional[int]` | `value` | | | diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index a44a6a866904..f459143ecfd7 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -74,6 +74,40 @@ scheduler_config: enable_prefix_aware_scheduling: false ``` +### Mamba Snapshot Boundaries + +Hybrid Mamba models must retain the recurrent Mamba state together with the +attention KV prefix. Snapshot policy is grouped under +`kv_cache_config.mamba_state_config`. `periodic_snapshot_interval` controls +periodic boundaries; set it to `0` to disable them. The interval is accepted +through this nested configuration in the Python API. YAML/JSON configuration +files also accept the deprecated `kv_cache_config.mamba_state_cache_interval` +key and migrate it to the nested field while loading. The prototype +`additional_snapshot_offsets_from_start` and +`additional_snapshot_offsets_from_end` options add fixed boundaries. Start +offsets count tokens from the beginning of the prompt. End offsets count +backward from the prompt end, and an end offset of `0` selects the final +prompt boundary. For example: + +```yaml +kv_cache_config: + enable_block_reuse: true + use_kv_cache_manager_v2: true + mamba_state_config: + periodic_snapshot_interval: 0 + additional_snapshot_offsets_from_start: [128] + additional_snapshot_offsets_from_end: [0, 32] +``` + +This retains snapshots after the first 128 tokens, at the end of the prompt, +and before the final 32 prompt tokens. Positions outside a particular prompt +are ignored. Exact explicit boundaries currently require aggregated serving +with `V2MambaHybridCacheManager`, `max_beam_width=1`, and no KV connector. +Hybrid Mamba models use the V1 C++ compatibility manager by default; select V2 +explicitly with `use_kv_cache_manager_v2: true`. V1 and the current +disaggregated-serving route support periodic snapshots only, while V2 does not +yet support disaggregated serving. + ### KV Cache Salting for Secure Reuse KV cache salting provides a security mechanism to control which requests can reuse cached KV states. When a `cache_salt` parameter is provided with a request, the KV cache system will only allow reuse of cached blocks given the same cache salt value. This prevents potential security issues such as prompt theft attacks, where malicious users might try to infer information from cached states of other users' requests. diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 0149ba7ae03a..032e8e1881fa 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -59,6 +59,7 @@ from .mamba_cache_manager import (BaseMambaCacheManager, CppMambaHybridCacheManager, MixedMambaHybridCacheManager, + V2MambaHybridCacheManager, use_py_mamba_cache_manager) from .model_engine import PyTorchModelEngine from .py_executor import PyExecutor @@ -94,14 +95,14 @@ def get_kv_cache_manager_cls( cache_transceiver_config: Optional[CacheTransceiverConfig] = None): """Resolve the concrete KV cache manager class for ``model_config``. - For hybrid mamba models the choice between ``Mixed`` (TRTLLM_USE_PY_MAMBA) - and ``Cpp`` (unified pool with block reuse) is made here. Callers that - don't care about disagg can omit ``is_disagg`` and get the unified-pool - default. + For hybrid mamba models the choice between + ``V2MambaHybridCacheManager`` and compatibility managers is made here. + Callers that don't care about disagg can omit ``is_disagg`` and get the + unified-pool default. - Env-var overrides (agg mode only — disagg picks its inner impl via - ``cache_transceiver_config.transceiver_runtime``): - * ``TRTLLM_USE_PY_MAMBA=1`` — Mixed manager with PythonMambaCacheManager. + V1 is the default for hybrid Mamba models. V2 is selected only by an + explicit ``kv_cache_config.use_kv_cache_manager_v2=True`` and is rejected + for disaggregated serving until its transceiver/page-table adapter exists. """ config = model_config.pretrained_config sparse_attn_config = model_config.sparse_attention_config @@ -122,44 +123,102 @@ def get_kv_cache_manager_cls( f"Sparse attention algorithm {sparse_attn_algorithm!r} is not " "supported with hybrid Mamba / linear-attention models.") + state_config = kv_cache_config.mamba_state_config + interval = state_config.periodic_snapshot_interval + has_periodic_snapshots = interval is not None and interval > 0 + has_additional_snapshots = bool( + state_config.additional_snapshot_offsets_from_start + or state_config.additional_snapshot_offsets_from_end) + use_v2 = kv_cache_config.use_kv_cache_manager_v2 is True + + if is_disagg and use_v2: + raise ValueError( + "KV cache manager V2 for hybrid Mamba models is not " + "supported with disaggregated serving. Set " + "use_kv_cache_manager_v2=False or 'auto' to use a V1 " + "Mamba cache manager.") + if has_additional_snapshots and not use_v2: + raise ValueError("Mamba additional snapshot offsets require " + "use_kv_cache_manager_v2=True; V1 supports only " + "periodic_snapshot_interval.") + if (kv_cache_config.enable_block_reuse and not has_periodic_snapshots + and not has_additional_snapshots): + raise ValueError( + "Hybrid Mamba block reuse requires at least one snapshot " + "policy: set " + "mamba_state_config.periodic_snapshot_interval > 0 or " + "provide additional snapshot offsets.") # Skip Softmax only changes attention kernels. Hybrid models still # need a Mamba-capable cache manager for recurrent state. + if is_disagg: + if kv_cache_config.enable_block_reuse: + return CppMambaHybridCacheManager + if (cache_transceiver_config is not None and + cache_transceiver_config.transceiver_runtime == "PYTHON"): + logger.info("Python transceiver detected; using " + "MixedMambaHybridCacheManager for hybrid model") + return MixedMambaHybridCacheManager + return CppMambaHybridCacheManager + if use_py_mamba_cache_manager(): + if use_v2: + raise ValueError( + "TRTLLM_USE_PY_MAMBA=1 conflicts with explicit " + "use_kv_cache_manager_v2=True.") if kv_cache_config.enable_block_reuse: raise ValueError( "TRTLLM_USE_PY_MAMBA=1 forces " "MixedMambaHybridCacheManager, which does not support " "block reuse. Disable block reuse or unset " - "TRTLLM_USE_PY_MAMBA to use CppMambaHybridCacheManager.") + "TRTLLM_USE_PY_MAMBA to use the configured cache manager.") logger.info( "Using MixedMambaHybridCacheManager for hybrid mamba model") return MixedMambaHybridCacheManager - if kv_cache_config.enable_block_reuse: - return CppMambaHybridCacheManager - if (cache_transceiver_config is not None - and cache_transceiver_config.transceiver_runtime == "PYTHON"): - logger.info("Python transceiver detected; using " - "MixedMambaHybridCacheManager for hybrid mamba model") - return MixedMambaHybridCacheManager - default_cls = CppMambaHybridCacheManager env_override = os.environ.get('TLLM_MAMBA_MANAGER_PREFERENCE', None) if env_override is not None: - if env_override.upper() == 'MIXED': + env_override = env_override.upper() + if env_override == 'MIXED': + if use_v2: + raise ValueError( + "TLLM_MAMBA_MANAGER_PREFERENCE=MIXED conflicts with " + "explicit use_kv_cache_manager_v2=True.") + if kv_cache_config.enable_block_reuse: + raise ValueError( + "TLLM_MAMBA_MANAGER_PREFERENCE=MIXED forces " + "MixedMambaHybridCacheManager, which does not support " + "block reuse. Disable block reuse, use the CPP " + "preference, or explicitly enable KV cache manager " + "V2.") logger.warning( - "Environment variable TLLM_MAMBA_MANAGER_PREFERENCE=MIXED overrides the default Mamba cache manager to MixedMambaHybridCacheManager. This may lead to increased memory usage due to lack of block reuse, but can be necessary for disaggregated setups or to avoid potential issues with the C++ manager. Set TLLM_MAMBA_MANAGER_PREFERENCE=CPP to use the CppMambaHybridCacheManager instead, which is the default for non-disaggregated setups without block reuse explicitly disabled." - ) + "Environment variable TLLM_MAMBA_MANAGER_PREFERENCE=MIXED " + "overrides the default Mamba cache manager to " + "MixedMambaHybridCacheManager.") return MixedMambaHybridCacheManager - elif env_override.upper() == 'CPP': + if env_override == 'CPP': + if use_v2: + raise ValueError( + "TLLM_MAMBA_MANAGER_PREFERENCE=CPP conflicts with " + "explicit use_kv_cache_manager_v2=True.") logger.warning( - "Environment variable TLLM_MAMBA_MANAGER_PREFERENCE=CPP overrides the default Mamba cache manager to CppMambaHybridCacheManager. This enables block reuse and can reduce memory usage, but may not be compatible with disaggregated setups. Set TLLM_MAMBA_MANAGER_PREFERENCE=MIXED to use the MixedMambaHybridCacheManager instead if you encounter issues with the C++ manager or are running in a disaggregated environment." - ) + "Environment variable TLLM_MAMBA_MANAGER_PREFERENCE=CPP " + "overrides the default Mamba cache manager to " + "CppMambaHybridCacheManager.") return CppMambaHybridCacheManager - else: - logger.warning( - f"Unrecognized value for TLLM_MAMBA_MANAGER_PREFERENCE: {env_override}. " - f"Expected 'CPP' or 'MIXED'. Using default {default_cls.__name__}." - ) - return default_cls + logger.warning( + f"Unrecognized value for TLLM_MAMBA_MANAGER_PREFERENCE: {env_override}. " + "Expected 'CPP' or 'MIXED'. Using the configured " + "KV cache manager default.") + + if not use_v2: + return CppMambaHybridCacheManager + + if (kv_cache_config.enable_block_reuse + and kv_cache_config.enable_kv_pool_rebalance): + raise ValueError( + "V2 Mamba block reuse is not compatible with " + "enable_kv_pool_rebalance because the rebalancer does not " + "yet model retained recurrent-state snapshots.") + return V2MambaHybridCacheManager elif sparse_attn_config is not None: return get_sparse_attn_kv_cache_manager(sparse_attn_config) else: @@ -358,22 +417,16 @@ def _get_model_kv_cache_manager_cls( cache_transceiver_config=self._cache_transceiver_config) cls = self._fallback_if_unsupported_kv_cache_manager_v2( cls, model_config, kv_cache_config) - # The V1-route hybrid mamba managers (disagg, TRTLLM_USE_CPP_MAMBA, - # TRTLLM_USE_PY_MAMBA, or one-model speculative decoding) keep mamba - # state in a separate cache that doesn't honor block reuse. Warn at - # the routing site so users see the warning where the decision is - # actually made. + # Compatibility managers do not support MTP block reuse. Warn at the + # routing site so users see the concrete manager selected for the + # incompatible combination. if is_hybrid_linear(model_engine.model.model_config.pretrained_config) \ - and kv_cache_config.enable_block_reuse: - uses_v1_mamba_route = self._is_disagg \ - or os.environ.get('TRTLLM_USE_CPP_MAMBA', '0') == '1' \ - or os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1' \ - or self._speculative_config is not None - if uses_v1_mamba_route: + and kv_cache_config.enable_block_reuse \ + and self._speculative_config is not None: + if not issubclass(cls, V2MambaHybridCacheManager): logger.warning( "Block reuse does not work with MTP for hybrid linear models " - "when using the legacy MambaCacheManager (TRTLLM_USE_CPP_MAMBA=1)" - ) + f"when using non-V2 Mamba cache manager {cls.__name__}") return cls def _fallback_if_unsupported_kv_cache_manager_v2( @@ -390,7 +443,7 @@ def _fallback_if_unsupported_kv_cache_manager_v2( if self._kv_connector_manager is not None: incompat.append("kv_connector_manager") if self._max_beam_width is not None and self._max_beam_width > 1: - incompat.append("beam_width > 1") + incompat.append("max_beam_width > 1") if incompat: incompat_str = ", ".join(incompat) # Some models are structurally bound to V2 and cannot fall @@ -415,6 +468,12 @@ def _fallback_if_unsupported_kv_cache_manager_v2( f"Gemma4 hybrid attention requires KVCacheManagerV2, " f"which is not yet supported with {incompat_str}. " f"Disable these features to run Gemma4 hybrid models.") + if is_hybrid_linear(config): + raise NotImplementedError( + "Hybrid Mamba cache managers do not support " + f"{incompat_str}; CppMambaHybridCacheManager does not " + "provide a compatible fallback. Use max_beam_width=1 " + "and disable the KV connector.") # Plain V2 (explicitly enabled or selected by a model default): # V2 was a preference, not a structural requirement, so we can # safely fall back to V1. @@ -446,6 +505,23 @@ def _per_manager_cache_cost(self, spec_config=self._speculative_config, **extra_kwargs)) + def _get_separate_target_layer_mask(self) -> Optional[List[bool]]: + """Return the target-only mask used by a separate draft cache.""" + if not self._should_create_separate_draft_kv_cache(): + return None + num_target_layers = (self._model_engine.model.model_config. + pretrained_config.num_hidden_layers) + return [True] * num_target_layers + + def _get_one_model_draft_layer_mask(self) -> List[bool]: + """Return the same draft-only mask used by runtime construction.""" + num_draft_layers = self._get_num_draft_layers() + if self._speculative_config.spec_dec_mode.is_external_drafter(): + return [True] * num_draft_layers + target_num_layers = (self._model_engine.model.model_config. + pretrained_config.num_hidden_layers) + return [False] * target_num_layers + [True] * num_draft_layers + def _get_kv_size_per_token(self, kv_cache_config: Optional[KvCacheConfig] = None ) -> CacheCost: @@ -457,8 +533,13 @@ def _get_kv_size_per_token(self, kv_cache_config = (kv_cache_config if kv_cache_config is not None else self._kv_cache_config) model_config = self._model_engine.model.model_config + target_layer_mask = self._get_separate_target_layer_mask() + target_kwargs = ({ + "layer_mask": target_layer_mask + } if target_layer_mask is not None else {}) total = self._per_manager_cache_cost(self._kv_cache_manager_cls, - model_config, kv_cache_config) + model_config, kv_cache_config, + **target_kwargs) if self._is_encoder_decoder(): total += CacheCost.from_raw(self._get_cross_kv_size_per_token()) if self._draft_model_engine is not None: @@ -493,7 +574,8 @@ def _get_kv_size_per_token(self, self._kv_cache_manager_cls, effective_draft_config, kv_cache_config, - num_layers=self._get_num_draft_layers()) + num_layers=self._get_num_draft_layers(), + layer_mask=self._get_one_model_draft_layer_mask()) return total def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction, @@ -1091,20 +1173,8 @@ def _create_one_model_draft_kv_cache_manager( Create a KV cache manager for draft model layers in one-model mode when target and draft have different KV cache layouts. """ - # Get target model's num_hidden_layers to compute correct layer indices. - # Draft model layers in one-model mode start at target_num_layers. - target_pretrained_config = self._model_engine.model.model_config.pretrained_config - target_num_layers = target_pretrained_config.num_hidden_layers - - # PARD, External Drafter: draft is a separate model, layers start from 0. - # Other methods (EAGLE3, MTP): draft layers are appended after target layers. num_draft_layers = self._get_num_draft_layers() - if self._speculative_config.spec_dec_mode.is_external_drafter(): - spec_dec_layer_mask = [True] * num_draft_layers - else: - spec_dec_layer_mask = [False] * target_num_layers + [ - True - ] * num_draft_layers + spec_dec_layer_mask = self._get_one_model_draft_layer_mask() # Get the effective draft config (explicit draft_config if available, # otherwise fall back to target model config for MTP). @@ -1191,9 +1261,13 @@ def _get_target_and_draft_cache_costs( target_kv_cache_config = (kv_cache_config if kv_cache_config is not None else self._kv_cache_config) total_kv = self._get_kv_size_per_token(target_kv_cache_config) + target_layer_mask = self._get_separate_target_layer_mask() + target_kwargs = ({ + "layer_mask": target_layer_mask + } if target_layer_mask is not None else {}) target_kv = self._per_manager_cache_cost( self._kv_cache_manager_cls, self._model_engine.model.model_config, - target_kv_cache_config) + target_kv_cache_config, **target_kwargs) # The draft contribution is whatever the aggregate has on top of the # target. Both pieces are CacheCost; subtraction is component-wise. draft_kv = CacheCost(slope=total_kv.slope - target_kv.slope, @@ -1941,6 +2015,9 @@ def _create_kv_cache_manager( mamba_ssm_stochastic_rounding = (stochastic_rounding and mamba_params.mamba_ssm_cache_dtype == torch.float16) + mamba_manager_extra_kwargs = dict(manager_extra_kwargs) + if not issubclass(kv_cache_manager_cls, V2MambaHybridCacheManager): + mamba_manager_extra_kwargs["model_type"] = "nemotron_hybrid" kv_cache_manager = kv_cache_manager_cls( # mamba cache parameters mamba_params.state_size, @@ -1968,10 +2045,9 @@ def _create_kv_cache_manager( spec_config=spec_config, is_estimating_kv_cache=estimating_kv_cache, execution_stream=execution_stream, - model_type="nemotron_hybrid", use_replay_state_update=use_replay, mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, - **manager_extra_kwargs, + **mamba_manager_extra_kwargs, ) elif is_qwen3_hybrid(config): if max_beam_width > 1: @@ -1988,7 +2064,6 @@ def _create_kv_cache_manager( spec_config=spec_config, quant_config=quant_config, ) - # Replay state update for GDN MTP: mirrors the nemotron_hybrid gating # above, minus the Mamba2-specific stochastic-rounding/Philox gate. # The GDN replay kernel does a plain cast on checkpoint commit, so @@ -2026,9 +2101,22 @@ def _create_kv_cache_manager( logger.info("GDN replay kernel is disabled; set " "TRTLLM_USE_GDN_REPLAY=1 to enable it") use_replay = False + + # Upstream GDN replay commits all local layer checkpoints through the + # contiguous C++ manager state view. V2 exposes per-layer state views, + # so enabling the same path there would fail for partitioned batches. + if (use_replay and issubclass(kv_cache_manager_cls, + V2MambaHybridCacheManager)): + logger.info( + "GDN replay is not supported by V2MambaHybridCacheManager; " + "using the legacy MTP path") + use_replay = False logger.info("GDN replay state update: " + ("ENABLED" if use_replay else "DISABLED")) + mamba_manager_extra_kwargs = dict(manager_extra_kwargs) + if not issubclass(kv_cache_manager_cls, V2MambaHybridCacheManager): + mamba_manager_extra_kwargs["model_type"] = "qwen3_next" kv_cache_manager = kv_cache_manager_cls( # mamba cache parameters mamba_params.state_size, @@ -2056,9 +2144,8 @@ def _create_kv_cache_manager( spec_config=spec_config, is_estimating_kv_cache=estimating_kv_cache, execution_stream=execution_stream, - model_type="qwen3_next", use_replay_state_update=use_replay, - **manager_extra_kwargs, + **mamba_manager_extra_kwargs, ) else: # NOTE: this is a workaround for VSWA to switch to calculate_max_num_blocks_for_vswa in KVCahceManager diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 84d20ce8caa5..8ef4967d1746 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -721,6 +721,7 @@ def __init__( execution_stream: Optional[torch.cuda.Stream] = None, is_disagg: bool = False, enable_stats: bool = False, + num_reserved_index_slots: int = 1, **kwargs, ) -> None: self.mapping = mapping @@ -1087,7 +1088,8 @@ def append_to_kv_heads_per_layer( # With pipeline parallelism, multiple microbatches can be in-flight # simultaneously, so we need slots for all concurrent sequences. - # Plus 1 for cuda graph dummy request. + # Reserve stable slots for persistent request IDs such as CUDA-graph + # padding requests. The default preserves the main-branch allocation. # In disaggregated mode, use a coefficient of 2: at any moment up to # `max_num_sequences` requests can be actively generating while another # up to `max_num_sequences` requests are still in KV transfer @@ -1095,10 +1097,16 @@ def append_to_kv_heads_per_layer( # capacity lets the next batch of active requests acquire slots without # waiting for the previous batch's transfers to finish. max_num_sequences = max_batch_size * mapping.pp_size - index_mapper_capacity = max_num_sequences * (2 if is_disagg else 1) + 1 + if num_reserved_index_slots < 0: + raise ValueError("num_reserved_index_slots must be non-negative") + index_mapper_capacity = ( + max_num_sequences * (2 if is_disagg else 1) + num_reserved_index_slots + ) logger.info( f"KVCacheManagerV2: IndexMapper capacity={index_mapper_capacity} " - f"(max_num_sequences={max_num_sequences}, is_disagg={is_disagg}, max_beam_width={max_beam_width})" + f"(max_num_sequences={max_num_sequences}, is_disagg={is_disagg}, " + f"num_reserved_index_slots={num_reserved_index_slots}, " + f"max_beam_width={max_beam_width})" ) self.index_mapper = IndexMapper(index_mapper_capacity, max_beam_width) self._early_freed_index_requests: set[int] = set() @@ -2698,7 +2706,10 @@ def get_iteration_stats(self): for window_size in sorted(windows) } - pool_group_ids = sorted(set(windows_by_pool_group) | set(pool_group_deltas)) + all_pool_group_ids = set(range(len(primary_stats))) + pool_group_ids = sorted( + all_pool_group_ids | set(windows_by_pool_group) | set(pool_group_deltas) + ) stats_by_pool_group = { pool_group_id: self._build_pool_group_iteration_stats( pool_group_id, diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index c413dd4fe20f..5357f709a649 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -17,7 +17,7 @@ import os from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Union +from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Tuple, Union import torch import triton @@ -27,19 +27,33 @@ from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( + BlockReusePolicy, KVCacheManagerV2, Role) from tensorrt_llm._torch.pyexecutor.llm_request import ( ATTENTION_DP_DUMMY_REQUEST_ID, LlmRequest) from tensorrt_llm._torch.pyexecutor.resource_manager import ( BaseResourceManager, CacheTypeCpp, DataType, KVCacheManager, PoolConfiguration, get_pp_layers) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests -from tensorrt_llm._utils import (nvtx_range, prefer_pinned, +from tensorrt_llm._utils import (TensorWrapper, convert_to_torch_tensor, + nvtx_range, prefer_pinned, torch_dtype_to_binding) from tensorrt_llm.bindings.internal.batch_manager import ( LinearAttentionMetadata, LinearCacheType) from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import (DEFAULT_BEAM_INDEX, + AttentionLayerConfig, + BatchDesc, BufferConfig, + CacheTierConfig, DataRole, + KVCacheDesc) +from tensorrt_llm.runtime.kv_cache_manager_v2 import \ + KVCacheManagerConfig as KVCacheManagerConfigPy +from tensorrt_llm.runtime.kv_cache_manager_v2 import (LayerId, PageIndexMode, + SsmLayerConfig, + SwaScratchReuseConfig, + exact_div) GB = 1 << 30 @@ -58,6 +72,56 @@ MIN_REPLAY_HISTORY_SIZE = 16 +def _get_num_cuda_graph_padding_dummy_slots( + spec_config: Optional["DecodingBaseConfig"], + max_batch_size: int, +) -> int: + """Return the number of persistent CUDA-graph padding dummy IDs.""" + if spec_config is None: + return 1 + + draft_len_schedule = getattr(spec_config, "draft_len_schedule", None) + spec_dec_mode = getattr(spec_config, "spec_dec_mode", None) + supports_dynamic_draft_len = (spec_dec_mode is not None and hasattr( + spec_dec_mode, "support_dynamic_draft_len") + and spec_dec_mode.support_dynamic_draft_len()) + if draft_len_schedule and supports_dynamic_draft_len: + runtime_draft_lengths = set() + first_uncovered_batch_size = 1 + for batch_size_threshold, draft_len in draft_len_schedule.items(): + if first_uncovered_batch_size > max_batch_size: + break + if batch_size_threshold >= first_uncovered_batch_size: + runtime_draft_lengths.add(draft_len) + first_uncovered_batch_size = batch_size_threshold + 1 + if first_uncovered_batch_size <= max_batch_size: + runtime_draft_lengths.add(0) + else: + max_draft_len = getattr(spec_config, "max_draft_len", 0) or 0 + max_total_draft_tokens = (getattr(spec_config, "max_total_draft_tokens", + 0) or 0) + is_linear_tree = getattr( + spec_config, + "is_linear_tree", + max_draft_len == max_total_draft_tokens, + ) + static_draft_len = (max_draft_len + if is_linear_tree else max_total_draft_tokens) + runtime_draft_lengths = {static_draft_len or 0} + + if ((getattr(spec_config, "acceptance_rate_window_size", 0) or 0) > 0 and + (getattr(spec_config, "acceptance_rate_threshold", 0) or 0) > 0): + runtime_draft_lengths.add(0) + return len(runtime_draft_lengths) + + +class MambaRole: + """V2 buffer roles owned only by the hybrid Mamba manager.""" + + SSM_STATE = DataRole("ssm_state") + CONV_STATE = DataRole("conv_state") + + def get_tensor_size_bytes(tensor): """Calculate tensor size in bytes.""" if isinstance(tensor, torch.Tensor): @@ -146,9 +210,9 @@ def use_py_mamba_cache_manager() -> bool: Returns True if TRTLLM_USE_PY_MAMBA='1' is set, False otherwise. Agg-mode-only override: forces the V1-route MixedMambaHybridCacheManager - with PythonMambaCacheManager inside instead of the default unified-pool - CppMambaHybridCacheManager. Disagg mode is unaffected — it already picks - PythonMambaCacheManager when transceiver_runtime='PYTHON'. + with PythonMambaCacheManager inside instead of the configured manager. + Disagg mode is unaffected — its compatibility routing selects Mixed or + Cpp based on the transceiver configuration. """ return os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1' @@ -996,7 +1060,8 @@ class MambaHybridCacheManager(BaseResourceManager, BaseMambaCacheManager): Used purely for ``isinstance`` / type-hint purposes so callers can refer to the family without caring about the concrete implementation. Concrete - selection (Mixed vs Cpp) lives in ``_util.py:_get_model_kv_cache_manager_cls``. + selection (V2, Mixed, or Cpp) lives in + ``_util.py:get_kv_cache_manager_cls``. """ @@ -1227,12 +1292,167 @@ def _promote_mamba_state_triton(dst: torch.Tensor, ) +def _mamba_snapshot_rule_counts( + kv_cache_config: KvCacheConfig, + max_seq_len: Optional[int], + tokens_per_block: int, +) -> Tuple[int, int]: + """Return reachable fixed rules and their partial-block upper bound.""" + if not kv_cache_config.enable_block_reuse: + return 0, 0 + + num_rules = 0 + num_unaligned_rules = 0 + state_config = kv_cache_config.mamba_state_config + for offset in set(state_config.additional_snapshot_offsets_from_start): + if max_seq_len is not None and offset > max_seq_len: + continue + num_rules += 1 + num_unaligned_rules += int(offset % tokens_per_block != 0) + for offset in set(state_config.additional_snapshot_offsets_from_end): + # A from-end offset is reachable iff some valid prompt is longer than + # the offset. Its absolute alignment depends on that prompt. + if max_seq_len is not None and offset >= max_seq_len: + continue + num_rules += 1 + num_unaligned_rules += 1 + return num_rules, num_unaligned_rules + + +def _mamba_regular_snapshot_interval( + kv_cache_config: KvCacheConfig, + max_seq_len: Optional[int], +) -> Optional[int]: + if not kv_cache_config.enable_block_reuse: + return None + interval = kv_cache_config.mamba_state_config.periodic_snapshot_interval + if interval is None or interval <= 0: + return None + if max_seq_len is not None and interval > max_seq_len: + return None + return interval + + +def _get_local_mamba_cache_layout( + model_config, + mapping: Mapping, + *, + spec_config=None, + layer_mask: Optional[List[bool]] = None, +): + """Return normalized params and local Mamba/attention layer counts. + + Cache construction and affine sizing must follow the model's PP layout: + partition base layers first, then place appended speculative layers on the + last PP rank. An explicit layer mask is authoritative for separate target + and draft caches. + """ + from tensorrt_llm._torch.pyexecutor.config_utils import \ + extract_mamba_kv_cache_params + + params = extract_mamba_kv_cache_params( + model_config.pretrained_config, + layer_mask=layer_mask, + spec_config=spec_config, + quant_config=model_config.quant_config, + ) + combined_layer_mask = [ + is_mamba or is_attention for is_mamba, is_attention in zip( + params.mamba_layer_mask, params.full_attention_layer_mask) + ] + local_layer_indices, _ = get_pp_layers( + sum(combined_layer_mask), + mapping, + spec_config=spec_config, + layer_mask=combined_layer_mask, + ) + local_mamba_layers = sum(params.mamba_layer_mask[layer_idx] + for layer_idx in local_layer_indices) + local_attention_layers = sum(params.full_attention_layer_mask[layer_idx] + for layer_idx in local_layer_indices) + return params, local_mamba_layers, local_attention_layers + + +def _estimate_mamba_hybrid_cache_cost( + model_config, + mapping: Mapping, + *, + max_batch_size: int, + kv_cache_config: KvCacheConfig, + num_layers: Optional[int], + tokens_per_block: int, + max_seq_len: Optional[int], + num_reserved_dummy_slots: int, + include_explicit_snapshots: bool, + cap_partial_attention_snapshots: bool, + **kwargs, +) -> Tuple[int, int]: + del num_layers + spec_config = kwargs.get("spec_config") + layer_mask = kwargs.get("layer_mask") + params, local_mamba_layers, local_attention_layers = ( + _get_local_mamba_cache_layout( + model_config, + mapping, + spec_config=spec_config, + layer_mask=layer_mask, + )) + attention_slope = (KVCacheManager.get_cache_size_per_token( + model_config, + mapping, + num_layers=local_attention_layers, + **kwargs, + ) if local_attention_layers > 0 else 0) + state_bytes_per_rank = (local_mamba_layers * + params.get_states_bytes_per_layer(mapping)) + max_resident_sequences = max_batch_size * mapping.pp_size + + if include_explicit_snapshots: + fixed_rules, unaligned_fixed_rules = _mamba_snapshot_rule_counts( + kv_cache_config, max_seq_len, tokens_per_block) + else: + fixed_rules = 0 + unaligned_fixed_rules = 0 + fixed_state_slots = (max_resident_sequences + num_reserved_dummy_slots + + max_resident_sequences * fixed_rules) + attention_block_bytes = attention_slope * tokens_per_block + + interval = _mamba_regular_snapshot_interval(kv_cache_config, max_seq_len) + has_unaligned_periodic_snapshot = (interval is not None + and interval % tokens_per_block != 0) + if cap_partial_attention_snapshots: + # V2 storage receives only an SSM slot count, which does not reveal + # whether each snapshot is block-aligned. Once any non-live SSM slot is + # possible, reserve one retained partial attention page per resident + # lineage. Fewer non-live slots and dummy slots only make this an + # overestimate; the bound does not rely on S >= 2N. + has_non_live_ssm_capacity = (fixed_rules > 0 or interval is not None + or num_reserved_dummy_slots > 0) + partial_attention_slots = (max_resident_sequences + if has_non_live_ssm_capacity else 0) + else: + partial_attention_slots = (max_resident_sequences * + unaligned_fixed_rules) + intercept = (fixed_state_slots * state_bytes_per_rank + + partial_attention_slots * attention_block_bytes) + + if interval is None: + regular_slope = 0 + else: + regular_slope = math.ceil(state_bytes_per_rank / interval) + if (has_unaligned_periodic_snapshot + and not cap_partial_attention_snapshots): + regular_slope += math.ceil(attention_block_bytes / interval) + return attention_slope + regular_slope, intercept + + class CppMambaHybridCacheManager(KVCacheManager, MambaHybridCacheManager): """Hybrid cache manager storing mamba states inside the KVCacheManager pool. Both KV cache blocks and recurrent state blocks are managed by the unified C++ KVCacheManager, enabling block reuse / prefix caching across attention - and mamba layers. This is the default hybrid manager. + and mamba layers. This compatibility manager remains available through the + manager preference override and legacy disaggregated routing. """ @@ -1351,7 +1571,7 @@ def __init__( self.kv_cache_config = kv_cache_config self.linear_attention_metadata = LinearAttentionMetadata() self.linear_attention_metadata.states_snapshot_interval = ( - kv_cache_config.mamba_state_cache_interval + kv_cache_config.mamba_state_config.periodic_snapshot_interval if kv_cache_config.enable_block_reuse else 0) return @@ -1411,7 +1631,7 @@ def __init__( self.linear_attention_metadata.cache_type = LinearCacheType.RECURRENT_STATES.value self.linear_attention_metadata.all_recurrent_states_bytes = self.ssm_bytes + self.conv_bytes self.linear_attention_metadata.states_snapshot_interval = ( - kv_cache_config.mamba_state_cache_interval + kv_cache_config.mamba_state_config.periodic_snapshot_interval if kv_cache_config.enable_block_reuse else 0) # RNN model params for disagg TP-mismatch split/concat. conv_section_map = {"nemotron_hybrid": 1, "qwen3_next": 2} @@ -1540,70 +1760,36 @@ def get_cache_size_per_token( max_batch_size: int, kv_cache_config: KvCacheConfig, num_layers: Optional[int] = None, + tokens_per_block: int = 32, + max_seq_len: Optional[int] = None, **kwargs, ): """Affine memory model for the unified hybrid KV pool. Returns ``(slope_bytes_per_token, intercept_bytes)``: - * ``slope`` = attention KV bytes per token (parent's formula) plus - the amortized regular-snapshot bytes per token from mamba layers. - * ``intercept`` = ``max_batch_size * num_mamba_layers_per_rank * - state_bytes_per_layer * STATIC_SLOTS_PER_REQUEST``. + * ``slope`` = attention KV bytes per token plus conservative periodic + Mamba-state and partial-attention-snapshot costs. + * ``intercept`` = live and CUDA-graph dummy Mamba state. Memory budget -> max tokens then becomes ``T = (budget - intercept) // slope`` instead of plain ``T = budget // bytes_per_token``. """ - # Lazy import to avoid pulling config_utils into module import order. - from tensorrt_llm._torch.pyexecutor.config_utils import \ - extract_mamba_kv_cache_params - - # Attention slope from the parent's existing formula. - attention_slope = KVCacheManager.get_cache_size_per_token( - model_config, mapping, num_layers=num_layers, **kwargs) - - params = extract_mamba_kv_cache_params( - model_config.pretrained_config, - quant_config=model_config.quant_config, + return _estimate_mamba_hybrid_cache_cost( + model_config, + mapping, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + num_layers=num_layers, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + num_reserved_dummy_slots=1, + include_explicit_snapshots=False, + cap_partial_attention_snapshots=False, + **kwargs, ) - state_bytes_per_layer = params.get_states_bytes_per_layer(mapping) - - # This not precise since pp layers are sharded by their order in model, not by their types. - # e.g. the upper half are all mamba layers while the lower half are all attention layers. - # But that's close enough for real world models where mamba and attention layers are interleaved - # and we don't have access with layer_masks at this point. - num_mamba_layers_per_rank = len( - mapping.pp_layers(params.num_mamba_layers)) - state_bytes_per_rank = num_mamba_layers_per_rank * state_bytes_per_layer - - # Per-request fixed cost. STATIC_SLOTS_PER_REQUEST = 1 today (the - # live mamba state); fixed-position snapshots are not yet - # implemented and would simply increment this constant. With - # pipeline parallelism, multiple microbatches are in-flight - # concurrently on the same rank, so each rank holds Mamba state - # for up to ``max_batch_size * pp_size`` concurrent sequences. - STATIC_SLOTS_PER_REQUEST = 1 - pp_size = mapping.pp_size if mapping is not None else 1 - intercept = (max_batch_size * pp_size * state_bytes_per_rank * - STATIC_SLOTS_PER_REQUEST) - - # Regular-snapshot bytes per token. None / non-positive intervals - # mean "no regular snapshots", so the mamba contribution is zero. - interval = kv_cache_config.mamba_state_cache_interval if kv_cache_config.enable_block_reuse else 0 - if interval is None or interval <= 0: - mamba_slope = 0 - else: - mamba_slope = state_bytes_per_rank // interval - # heuristic: When block reuse is enabled, we assume the mamba snapshots are dominant instead of active states, - # otherwise we may run out of kv cache blocks prior to mamba blocks due to the large number of max_batch_size. - # So we ignore intercept and only calculate max_tokens based on slope - # This can be improved by a more accurate max_batch_size and ISL/OSL estimation in the future. - if mamba_slope > 0: - intercept = 0 - return attention_slope + mamba_slope, intercept - @property def use_gdn_cached_replay_all_layer_commit(self) -> bool: return self._use_gdn_cached_replay_all_layer_commit @@ -2044,6 +2230,11 @@ def _setup_state_indices(self, requests=None) -> None: def get_state_indices(self, request_ids: Optional[List[int]] = None, is_padding: Optional[List[bool]] = None) -> list: + if self.local_num_mamba_layers == 0: + # Mamba metadata is prepared on every PP rank even when this rank + # owns only attention layers. No local kernel consumes these + # indices, so avoid consulting state that is intentionally absent. + return [0] * len(request_ids) if request_ids is not None else [] if request_ids is not None: # Return indices in the order of the caller's request_ids, # not the internal self.requests order. This is critical when @@ -2261,3 +2452,1360 @@ def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: """Return the persistent (cache_size,) int64 Philox seed buffer or None when stochastic rounding is not active for this manager.""" return getattr(self, 'mamba_ssm_rand_seed', None) + + +class V2MambaHybridCacheManager(KVCacheManagerV2, MambaHybridCacheManager): + """Hybrid Mamba cache manager backed by KVCacheManagerV2. + + Attention KV pages and Mamba recurrent-state pages are both owned by the + Python V2 cache manager. Mamba layers are represented as V2 SSM layers, + while this wrapper exposes the state tensors and slot indices expected by + the PyTorch Mamba kernels. + """ + + def __init__( + self, + # mamba cache parameters + mamba_d_state: int, + mamba_d_conv: int, + mamba_num_heads: int, + mamba_n_groups: int, + mamba_head_dim: int, + mamba_num_layers: int, + mamba_layer_mask: List[bool], + mamba_cache_dtype: torch.dtype, + mamba_ssm_cache_dtype: torch.dtype, + kv_cache_config: KvCacheConfig, + kv_cache_type: CacheTypeCpp, + *, + num_layers: int, + num_kv_heads: Union[int, List[Optional[int]]], + head_dim: int, + tokens_per_block: int, + max_seq_len: int, + max_batch_size: int, + mapping: Mapping, + dtype: DataType = DataType.HALF, + spec_config: Optional["DecodingBaseConfig"] = None, + layer_mask: Optional[List[bool]] = None, + is_estimating_kv_cache: bool = False, + is_draft: bool = False, + use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, + **kwargs, + ) -> None: + total_layers = len(mamba_layer_mask) + if layer_mask is None: + full_attention_layer_mask = [False] * total_layers + elif len(layer_mask) != total_layers: + raise ValueError( + f"layer_mask length ({len(layer_mask)}) must match " + f"mamba_layer_mask length ({total_layers})") + else: + full_attention_layer_mask = list(layer_mask) + + combined_layer_mask = [ + mamba_layer_mask[i] or full_attention_layer_mask[i] + for i in range(total_layers) + ] + + self._mamba_layer_mask = list(mamba_layer_mask) + self.requests = [] + self._use_replay_state_update = use_replay_state_update + self.replay_step_width: Optional[int] = ( + spec_config.tokens_per_gen_step + if spec_config is not None and use_replay_state_update else None) + self.replay_history_size: Optional[int] = self.replay_step_width + self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding + self._seed_rank_offset = _mamba_rank_offset(mapping) + self._seed_request_counter = 0 + self._num_cuda_graph_padding_dummy_slots = ( + _get_num_cuda_graph_padding_dummy_slots(spec_config, + max_batch_size)) + self._num_attention_dp_dummy_slots = int(mapping.enable_attention_dp) + self._num_reserved_dummy_slots = ( + self._num_cuda_graph_padding_dummy_slots + + self._num_attention_dp_dummy_slots) + self.ssm_state_dtype = (mamba_ssm_cache_dtype if mamba_ssm_cache_dtype + is not None else mamba_cache_dtype) + self.conv_state_dtype = mamba_cache_dtype + + self.pp_layers, _ = get_pp_layers( + mamba_num_layers + num_layers, + mapping, + spec_config=spec_config, + layer_mask=combined_layer_mask, + ) + self.mamba_pp_layers = [ + layer_idx for layer_idx in self.pp_layers + if mamba_layer_mask[layer_idx] + ] + self.local_num_mamba_layers = len(self.mamba_pp_layers) + + if self.local_num_mamba_layers > 0: + tp_size = mapping.tp_size if not mapping.enable_attention_dp else 1 + d_inner = mamba_head_dim * mamba_num_heads + conv_dim = d_inner + 2 * mamba_n_groups * mamba_d_state + nheads = mamba_num_heads + assert nheads % tp_size == 0, "mamba_num_heads must be divisible by tp_size" + assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" + if use_replay_state_update: + assert mamba_n_groups % tp_size == 0, \ + "replay state update requires mamba_n_groups divisible by tp_size" + self._n_groups_per_rank = mamba_n_groups // tp_size + conv_dim = conv_dim // tp_size + nheads = nheads // tp_size + self.conv_state_shape = [conv_dim, mamba_d_conv - 1] + self.ssm_state_shape = [nheads, mamba_head_dim, mamba_d_state] + self.ssm_count = math.prod(self.ssm_state_shape) + self.conv_count = math.prod(self.conv_state_shape) + self.ssm_bytes = self.ssm_count * self.ssm_state_dtype.itemsize + self.conv_bytes = self.conv_count * self.conv_state_dtype.itemsize + else: + logger.info( + "No local mamba layers for this rank, skipping mamba state views" + ) + self._n_groups_per_rank = 0 + self.conv_state_shape = [] + self.ssm_state_shape = [] + self.ssm_count = 0 + self.conv_count = 0 + self.ssm_bytes = 0 + self.conv_bytes = 0 + + if isinstance(num_kv_heads, int): + per_layer_kv_heads = [num_kv_heads] * total_layers + else: + if len(num_kv_heads) != total_layers: + raise ValueError( + f"num_kv_heads list length ({len(num_kv_heads)}) does not " + f"match total layers ({total_layers})") + per_layer_kv_heads = list(num_kv_heads) + for i, is_mamba in enumerate(mamba_layer_mask): + if is_mamba: + per_layer_kv_heads[i] = 0 + + self._setup_mtp_intermediate_states(spec_config, max_batch_size) + + kv_cache_config = kv_cache_config.model_copy(deep=True) + if any(mamba_layer_mask) and kv_cache_config.enable_block_reuse: + # SSM reuse is valid only at explicit snapshot boundaries. + kv_cache_config.block_reuse_policy = BlockReusePolicy.PER_REQUEST.value + self.kv_cache_config = kv_cache_config + + super().__init__( + kv_cache_config, + kv_cache_type, + num_layers=mamba_num_layers + num_layers, + num_kv_heads=per_layer_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + mapping=mapping, + dtype=dtype, + spec_config=spec_config, + layer_mask=combined_layer_mask, + is_draft=is_draft, + is_estimating_kv_cache=is_estimating_kv_cache, + num_reserved_index_slots=self._num_reserved_dummy_slots, + **kwargs, + ) + + self.mamba_layer_offsets = { + layer_id: offset + for offset, layer_id in enumerate(self.mamba_pp_layers) + } + self.mamba_local_layer_ids = [ + self.layer_offsets[layer_id] for layer_id in self.mamba_pp_layers + ] + self._request_id_to_state_index = {} + + state_index_capacity = (self.max_batch_size + + self._num_reserved_dummy_slots) + self.cuda_state_indices = torch.zeros([state_index_capacity], + dtype=torch.int32, + device="cuda") + self._host_state_indices = torch.zeros([state_index_capacity], + dtype=torch.int32, + pin_memory=prefer_pinned()) + + if self.local_num_mamba_layers > 0: + first_mamba_local_layer = self.mamba_local_layer_ids[0] + self.ssm_layer_group_id = self.impl.get_layer_group_id( + LayerId(first_mamba_local_layer)) + self._ssm_page_index_scale = self.impl.get_page_index_scale( + LayerId(first_mamba_local_layer), MambaRole.SSM_STATE) + num_ssm_pages = self.impl.get_page_index_upper_bound( + LayerId(first_mamba_local_layer), MambaRole.SSM_STATE) + num_ssm_slots = ((num_ssm_pages + self._ssm_page_index_scale - 1) // + self._ssm_page_index_scale) + required_live_slots = (self._max_resident_sequences() + + self._num_reserved_dummy_slots) + if num_ssm_slots < required_live_slots: + KVCacheManagerV2.shutdown(self) + raise ValueError( + "The V2 Mamba state pool has only " + f"{num_ssm_slots} slots but needs at least " + f"{required_live_slots} live/dummy slots. Increase the " + "KV cache budget or allocate a larger Mamba pool_ratio.") + self._setup_states() + self._setup_replay_buffers(spec_config) + else: + self.ssm_layer_group_id = None + self._ssm_page_index_scale = 1 + self.all_ssm_states = [] + self.all_conv_states = [] + self._setup_replay_buffers(spec_config) + + @staticmethod + def get_cache_size_per_token(model_config, + mapping: Mapping, + *, + max_batch_size: int, + kv_cache_config: KvCacheConfig, + num_layers: Optional[int] = None, + tokens_per_block: int = 32, + max_seq_len: Optional[int] = None, + **kwargs): + spec_config = kwargs.get("spec_config") + num_reserved_dummy_slots = (_get_num_cuda_graph_padding_dummy_slots( + spec_config, max_batch_size) + int(mapping.enable_attention_dp)) + return _estimate_mamba_hybrid_cache_cost( + model_config, + mapping, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + num_layers=num_layers, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + num_reserved_dummy_slots=num_reserved_dummy_slots, + include_explicit_snapshots=True, + cap_partial_attention_snapshots=True, + **kwargs, + ) + + def _is_local_mamba_layer(self, local_layer_idx: int) -> bool: + return self._mamba_layer_mask[self.pp_layers[local_layer_idx]] + + def _get_pool_page_index_role(self, pool_id: int) -> DataRole: + layer_id = int(self.impl.layer_grouping[pool_id][0]) + if self._is_local_mamba_layer(layer_id): + return MambaRole.SSM_STATE + return Role.KEY + + def _get_pool_paired_role(self, pool_id: int) -> Optional[DataRole]: + layer_id = int(self.impl.layer_grouping[pool_id][0]) + if self._is_local_mamba_layer(layer_id): + return None + return (Role.VALUE + if self.kv_cache_type != CacheTypeCpp.SELFKONLY else None) + + def _get_block_scale_role_for_pool(self, + pool_id: int) -> Optional[DataRole]: + if (self.dtype != DataType.NVFP4 + or self._get_pool_page_index_role(pool_id) != Role.KEY): + return None + return Role.KEY_BLOCK_SCALE + + def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: + """Build attention-compatible tables for mixed attention/state pools.""" + kv_cache_pool_pointers_list = [] + kv_cache_pool_mapping_list = [] + block_scale_pool_pointers_list = [] + if self.enable_swa_scratch_reuse: + for local_layer_idx in range(self.num_local_layers): + layer_id = LayerId(local_layer_idx) + pool_id = self.impl.get_layer_group_id(layer_id) + page_index_role = self._get_pool_page_index_role(pool_id) + kv_cache_pool_pointers_list.append([ + self.impl.get_mem_pool_base_address( + layer_id, page_index_role, PageIndexMode.PER_LAYER), + 0, + ]) + if self.dtype == DataType.NVFP4: + block_scale_role = self._get_block_scale_role_for_pool( + pool_id) + block_scale_pool_pointers_list.append([ + self.impl.get_mem_pool_base_address( + layer_id, block_scale_role, PageIndexMode.PER_LAYER) + if block_scale_role is not None else 0, + 0, + ]) + kv_cache_pool_mapping_list.append([int(layer_id), 0]) + else: + for pool_id in range(self.num_pools): + layer_id = self.impl.layer_grouping[pool_id][0] + page_index_role = self._get_pool_page_index_role(pool_id) + kv_cache_pool_pointers_list.append([ + self.impl.get_mem_pool_base_address(layer_id, + page_index_role, + PageIndexMode.SHARED), + 0, + ]) + if self.dtype == DataType.NVFP4: + block_scale_role = self._get_block_scale_role_for_pool( + pool_id) + block_scale_pool_pointers_list.append([ + self.impl.get_mem_pool_base_address( + layer_id, block_scale_role, PageIndexMode.SHARED) + if block_scale_role is not None else 0, + 0, + ]) + + for local_layer_idx in range(self.num_local_layers): + layer_id = LayerId(local_layer_idx) + layer_group_id = self.impl.get_layer_group_id(layer_id) + page_index_role = self._get_pool_page_index_role(layer_group_id) + index_base_addr = kv_cache_pool_pointers_list[layer_group_id][0] + addr_offset = (self.impl.get_mem_pool_base_address( + layer_id, page_index_role, PageIndexMode.SHARED) - + index_base_addr) + offset_divisor = self.impl.get_page_stride( + layer_id, page_index_role) + if self._get_pool_paired_role(layer_group_id) is not None: + offset_divisor *= self.kv_factor + offset = exact_div(addr_offset, offset_divisor) + + if self.dtype != DataType.NVFP4: + block_scale_offset = None + else: + block_scale_role = self._get_block_scale_role_for_pool( + layer_group_id) + if block_scale_role is None: + block_scale_offset = None + else: + block_scale_base_addr = block_scale_pool_pointers_list[ + layer_group_id][0] + block_scale_addr_offset = ( + self.impl.get_mem_pool_base_address( + layer_id, block_scale_role, + PageIndexMode.SHARED) - block_scale_base_addr) + block_scale_divisor = self.impl.get_page_stride( + layer_id, block_scale_role) + if self._get_pool_paired_role( + layer_group_id) is not None: + block_scale_divisor *= self.kv_factor + block_scale_offset = exact_div(block_scale_addr_offset, + block_scale_divisor) + + if block_scale_offset is not None: + assert block_scale_offset == offset, ( + "Block scale offset and offset should be the same") + + kv_cache_pool_mapping_list.append([layer_group_id, offset]) + + if self.dtype == DataType.NVFP4: + for pool_id, block_scale_pool_pointers in enumerate( + block_scale_pool_pointers_list): + pool_pointers = kv_cache_pool_pointers_list[pool_id] + kv_cache_pool_pointers_list[pool_id] = [ + [pool_pointers[0], block_scale_pool_pointers[0]], + [pool_pointers[1], block_scale_pool_pointers[1]], + ] + + self.kv_cache_pool_pointers = torch.tensor( + kv_cache_pool_pointers_list, + dtype=torch.int64, + device="cpu", + pin_memory=prefer_pinned(), + ) + self.kv_cache_pool_mapping = torch.tensor( + kv_cache_pool_mapping_list, + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), + ) + self.index_scales = torch.empty( + self.num_pools, + dtype=torch.int32, + pin_memory=prefer_pinned(), + device="cpu", + ) + self.kv_offset = torch.empty( + self.num_pools, + dtype=torch.int32, + pin_memory=prefer_pinned(), + device="cpu", + ) + for pool_id in range(self.num_pools): + layer_id = self.impl.layer_grouping[pool_id][0] + page_index_role = self._get_pool_page_index_role(pool_id) + self.index_scales[pool_id] = self.impl.get_page_index_scale( + layer_id, page_index_role) + paired_role = self._get_pool_paired_role(pool_id) + if paired_role is not None: + self.kv_offset[pool_id] = exact_div( + self.impl.get_mem_pool_base_address(layer_id, paired_role, + PageIndexMode.SHARED) - + self.impl.get_mem_pool_base_address( + layer_id, page_index_role, PageIndexMode.SHARED), + self.impl.get_page_stride(layer_id, page_index_role), + ) + else: + self.kv_offset[pool_id] = 0 + self._index_scale_ints = self.index_scales.tolist() + + self.host_kv_cache_block_offsets = torch.zeros( + self.num_pools, + index_mapper_capacity * self.max_beam_width, + 2, + self.max_blocks_per_seq, + dtype=torch.int32, + pin_memory=prefer_pinned(), + device="cpu", + ) + if self.enable_swa_scratch_reuse: + self._prepare_swa_scratch_copy_tensors(index_mapper_capacity) + + def _prepare_swa_scratch_copy_tensors(self, + index_mapper_capacity: int) -> None: + pool_ids = torch.empty( + self.num_attention_op_pools, + 2, + dtype=torch.long, + device="cpu", + ) + scales = torch.empty( + self.num_attention_op_pools, + 2, + dtype=torch.int32, + device="cpu", + ) + layer_offsets = torch.empty_like(scales) + scratch_pages = torch.empty_like(scales) + for local_layer_idx in range(self.num_local_layers): + layer_id = LayerId(local_layer_idx) + pool_id = self.layer_to_pool_mapping_dict[layer_id] + page_index_role = self._get_pool_page_index_role(pool_id) + paired_role = self._get_pool_paired_role(pool_id) + roles = [page_index_role, paired_role or page_index_role] + for role_idx, role in enumerate(roles): + converter = self.impl.get_page_index_converter(layer_id, role) + if converter.expansion != 1: + raise NotImplementedError( + "SWA scratch block-table conversion does not support " + "expanded page indices yet: " + f"layer={layer_id}, role={role}, " + f"expansion={converter.expansion}") + pool_ids[local_layer_idx, role_idx] = pool_id + scales[local_layer_idx, role_idx] = int(converter.scale) + layer_offsets[local_layer_idx, + role_idx] = int(converter.layer_offset) + scratch_pages[local_layer_idx, + role_idx] = int(converter.scratch_pages_per_block) + + staging_capacity = index_mapper_capacity * self.max_beam_width + device = torch.device("cuda", torch.cuda.current_device()) + self._device_kv_cache_block_offsets_input = torch.empty_like( + self.host_kv_cache_block_offsets, device=device) + self._device_attention_op_block_offsets_staging = torch.empty( + self.num_attention_op_pools, + staging_capacity, + 2, + self.max_blocks_per_seq, + dtype=torch.int32, + device=device, + ) + self._device_copy_idx_staging = torch.zeros(staging_capacity, + dtype=torch.long, + device=device) + self._device_num_contexts = torch.empty((), + dtype=torch.int32, + device=device) + self._device_attention_op_pool_ids = pool_ids.to(device=device) + self._device_attention_op_scales = scales.to(device=device) + self._device_attention_op_layer_offsets = layer_offsets.to( + device=device) + self._device_attention_op_scratch_pages = scratch_pages.to( + device=device) + self._device_block_positions = torch.arange(self.max_blocks_per_seq, + dtype=torch.int32, + device=device) + + min_scale = int(scales.min().item()) if scales.numel() > 0 else 1 + max_scratch_pages = (int(scratch_pages.max().item()) + if scratch_pages.numel() > 0 else 1) + self._max_scratch_slots = max( + 1, + (self.max_blocks_per_seq * max_scratch_pages + min_scale - 1) // + min_scale, + ) + scratch_slots_shape = (self.num_pools, staging_capacity, + self._max_scratch_slots) + self._host_scratch_begs_staging = torch.zeros( + self.num_pools, + staging_capacity, + dtype=torch.int32, + pin_memory=prefer_pinned(), + device="cpu", + ) + self._host_scratch_ends_staging = torch.zeros( + self.num_pools, + staging_capacity, + dtype=torch.int32, + pin_memory=prefer_pinned(), + device="cpu", + ) + self._host_scratch_slots_staging = torch.zeros( + scratch_slots_shape, + dtype=torch.int32, + pin_memory=prefer_pinned(), + device="cpu", + ) + self._device_scratch_begs_staging = torch.zeros( + self.num_pools, + staging_capacity, + dtype=torch.int32, + device=device, + ) + self._device_scratch_ends_staging = torch.zeros( + self.num_pools, + staging_capacity, + dtype=torch.int32, + device=device, + ) + self._device_scratch_slots_staging = torch.zeros( + scratch_slots_shape, + dtype=torch.int32, + device=device, + ) + + def _get_event_window_sizes_by_layer_group(self) -> Dict[int, int]: + + def get_event_window_size(layer_id: int) -> int: + layer_config = self.kv_cache_manager_py_config.layers[layer_id] + window_size = getattr(layer_config, "sliding_window_size", None) + return (self.max_seq_len + if window_size is None else int(window_size)) + + return { + int(layer_group_id): get_event_window_size(int(layer_ids[0])) + for layer_group_id, layer_ids in enumerate(self.impl.layer_grouping) + } + + def _max_resident_sequences(self) -> int: + return self.max_batch_size * self.mapping.pp_size + + def _mamba_state_bytes_per_slot(self) -> int: + return self.local_num_mamba_layers * (self.ssm_bytes + self.conv_bytes) + + def _num_ssm_snapshots_for_capacity( + self, + capacity: int, + kv_cache_config: KvCacheConfig, + ) -> int: + if capacity <= 0 or not kv_cache_config.enable_block_reuse: + return 0 + + fixed_rules, _ = _mamba_snapshot_rule_counts(kv_cache_config, + self.max_seq_len, + self.tokens_per_block) + interval = _mamba_regular_snapshot_interval(kv_cache_config, + self.max_seq_len) + regular_snapshots = capacity // interval if interval is not None else 0 + return (self._max_resident_sequences() * fixed_rules + + regular_snapshots) + + def _ssm_slots_per_request_for_typical_batch( + self, + capacity: int, + kv_cache_config: KvCacheConfig, + ) -> List[int]: + num_sequences = self._max_resident_sequences() + fixed_rules = 0 + if capacity > 0: + fixed_rules, _ = _mamba_snapshot_rule_counts( + kv_cache_config, self.max_seq_len, self.tokens_per_block) + interval = _mamba_regular_snapshot_interval(kv_cache_config, + self.max_seq_len) + regular_snapshots = capacity // interval if interval is not None else 0 + regular_per_request, regular_remainder = divmod(regular_snapshots, + num_sequences) + dummy_per_request, dummy_remainder = divmod( + self._num_reserved_dummy_slots, num_sequences) + return [ + 1 + fixed_rules + regular_per_request + int(i < regular_remainder) + + dummy_per_request + int(i < dummy_remainder) + for i in range(num_sequences) + ] + + def _get_quota_from_max_tokens(self, max_tokens: int) -> int: + attention_quota = super()._get_quota_from_max_tokens(max_tokens) + ssm_slots = self._ssm_slots_per_request_for_typical_batch( + max_tokens, self.kv_cache_config) + state_slots = sum(ssm_slots) + state_quota = state_slots * self._mamba_state_bytes_per_slot() + num_request_lineages = len(ssm_slots) + # Once the plan contains any non-live SSM capacity, reserve one partial + # attention page per request lineage. This remains conservative when + # the plan contains fewer than one non-live slot per lineage. + extra_attention_quota = (num_request_lineages * + self._attention_cache_bytes_per_token() * + self.tokens_per_block + if state_slots > num_request_lineages else 0) + return attention_quota + state_quota + extra_attention_quota + + def _get_max_tokens_from_quota(self, quota: int) -> float: + if self._get_quota_from_max_tokens(0) > quota: + return 0 + + low = 0 + high = 1 + while self._get_quota_from_max_tokens(high) <= quota: + low = high + high *= 2 + if high >= 1 << 62: + return float("inf") + + while low + 1 < high: + mid = (low + high) // 2 + if self._get_quota_from_max_tokens(mid) <= quota: + low = mid + else: + high = mid + return low + + def _planned_token_capacity( + self, + kv_cache_config: KvCacheConfig, + gpu_quota: int, + ) -> int: + capacity = self._get_max_tokens_from_quota(gpu_quota) + if math.isinf(capacity): + capacity = (kv_cache_config.max_tokens if kv_cache_config.max_tokens + is not None else self.max_seq_len * + self._max_resident_sequences()) + capacity = min( + capacity, + self.max_seq_len * self._max_resident_sequences(), + ) + if kv_cache_config.max_tokens is not None: + capacity = min(capacity, kv_cache_config.max_tokens) + return max(0, int(capacity)) + + def _minimum_live_gpu_quota(self) -> int: + """Return the minimum quota for live states and one attention page.""" + num_sequences = self._max_resident_sequences() + state_slots = num_sequences + self._num_reserved_dummy_slots + state_quota = state_slots * self._mamba_state_bytes_per_slot() + attention_block_quota = (self._attention_cache_bytes_per_token() * + self.tokens_per_block) + # At zero token capacity, any reserved dummy state makes non-live SSM + # capacity possible. Match the normal estimator by reserving one + # partial attention page per request lineage in that case. + extra_attention_quota = (num_sequences * attention_block_quota + if state_slots > num_sequences else 0) + attention_quota = max( + super()._get_quota_from_max_tokens(0) + extra_attention_quota, + attention_block_quota, + ) + return state_quota + attention_quota + + def _build_cache_config( + self, + kv_cache_config: KvCacheConfig, + *, + tokens_per_block: int, + vocab_size: Optional[int], + cache_tiers: List[CacheTierConfig], + ): + # Kept in the virtual method contract for cache-manager subclasses. + # The generic V2 config no longer stores the vocabulary size. + del vocab_size + gpu_quota = cache_tiers[0].quota + minimum_live_quota = self._minimum_live_gpu_quota() + if minimum_live_quota > gpu_quota: + raise ValueError( + "The V2 Mamba GPU cache quota is too small for live recurrent " + f"states and attention pages: got {gpu_quota} bytes, need at " + f"least {minimum_live_quota} bytes.") + buffer_type = [Role.KEY] + if self.kv_cache_type != CacheTypeCpp.SELFKONLY: + buffer_type.append(Role.VALUE) + if kv_cache_config.dtype == "nvfp4": + for layer_idx, hd in enumerate(self.head_dim_per_layer): + if self._is_local_mamba_layer(layer_idx): + continue + assert hd % 2 == 0, ( + f"head_dim must be divisible by 2 for nvfp4 kv cache, " + f"but layer {layer_idx} has head_dim={hd}") + buffer_type.append(Role.KEY_BLOCK_SCALE) + if self.kv_cache_type != CacheTypeCpp.SELFKONLY: + buffer_type.append(Role.VALUE_BLOCK_SCALE) + + scratch_reuse_config = None + if self.enable_swa_scratch_reuse: + scratch_reuse_config = SwaScratchReuseConfig( + max_rewind_len=self.num_extra_kv_tokens) + + layers = [] + for local_layer_idx, global_layer_idx in enumerate(self.pp_layers): + layer_id = LayerId(local_layer_idx) + if self._mamba_layer_mask[global_layer_idx]: + layers.append( + SsmLayerConfig( + layer_id=layer_id, + buffers=[ + BufferConfig(role=MambaRole.SSM_STATE, + size=self.ssm_bytes), + BufferConfig(role=MambaRole.CONV_STATE, + size=self.conv_bytes), + ], + )) + else: + layers.append( + AttentionLayerConfig( + layer_id=layer_id, + buffers=[ + BufferConfig( + role=role, + size=self.get_layer_bytes_per_token( + local_layer_idx=layer_id, data_role=role) * + tokens_per_block, + ) for role in buffer_type + ], + sliding_window_size=self.max_attention_window_vec[ + global_layer_idx % + len(self.max_attention_window_vec)], + num_sink_tokens=None, + )) + + num_sequences = self._max_resident_sequences() + planned_capacity = self._planned_token_capacity(kv_cache_config, + cache_tiers[0].quota) + capacity_per_request, capacity_remainder = divmod( + planned_capacity, num_sequences) + typical_ssm_slots = self._ssm_slots_per_request_for_typical_batch( + planned_capacity, kv_cache_config) + + return KVCacheManagerConfigPy( + tokens_per_block=tokens_per_block, + cache_tiers=cache_tiers, + max_util_for_resume=kv_cache_config.max_util_for_resume, + initial_pool_ratio=kv_cache_config.pool_ratio, + layers=layers, + enable_partial_reuse=kv_cache_config.enable_partial_reuse, + typical_step=BatchDesc([ + KVCacheDesc( + capacity=capacity_per_request + int(i < capacity_remainder), + history_length=max( + 0, + capacity_per_request + int(i < capacity_remainder) - 1), + num_ssm_slots=typical_ssm_slots[i], + ) for i in range(num_sequences) + ]), + swa_scratch_reuse=scratch_reuse_config, + # SSM lifecycles require minimum-snapshot commit semantics. The + # flag is harmless when reuse is disabled because no commits are + # attempted, while the runtime config still needs the invariant. + commit_min_snapshot=True, + enable_stats=self.enable_stats, + ) + + def _get_state_buffer(self, local_layer_idx: int, role, dtype: torch.dtype, + state_shape: List[int]) -> torch.Tensor: + addr = self.impl.get_mem_pool_base_address(LayerId(local_layer_idx), + role, PageIndexMode.SHARED) + num_pages = self.impl.get_page_index_upper_bound( + LayerId(local_layer_idx), role) + raw = convert_to_torch_tensor( + TensorWrapper(addr, dtype, [num_pages] + state_shape)) + page_index_scale = self.impl.get_page_index_scale( + LayerId(local_layer_idx), role) + num_slots = (num_pages + page_index_scale - 1) // page_index_scale + # V2 coalesces same-size per-layer buffers inside each slot. Kernels + # index Mamba states by logical slot id, so expose only this layer's + # sub-page from each coalesced slot instead of the raw page-index view. + return raw.as_strided( + [num_slots] + state_shape, + [raw.stride(0) * page_index_scale] + list(raw.stride()[1:]), + ) + + def _setup_states(self) -> None: + self.all_ssm_states = [ + self._get_state_buffer(local_layer_idx, MambaRole.SSM_STATE, + self.ssm_state_dtype, self.ssm_state_shape) + for local_layer_idx in self.mamba_local_layer_ids + ] + self.all_conv_states = [ + self._get_state_buffer(local_layer_idx, MambaRole.CONV_STATE, + self.conv_state_dtype, self.conv_state_shape) + for local_layer_idx in self.mamba_local_layer_ids + ] + + def _setup_mtp_intermediate_states(self, spec_config, + max_batch_size) -> None: + self.spec_config = spec_config + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None + if self.spec_config is not None and self.local_num_mamba_layers > 0: + speculative_num_draft_tokens = self.spec_config.tokens_per_gen_step - 1 + if not self._use_replay_state_update: + self.intermediate_ssm_states = torch.zeros( + size=[ + self.local_num_mamba_layers, max_batch_size, + speculative_num_draft_tokens + 1 + ] + self.ssm_state_shape, + dtype=self.ssm_state_dtype, + device="cuda", + ) + self.intermediate_conv_states = torch.zeros( + size=[ + self.local_num_mamba_layers, max_batch_size, + speculative_num_draft_tokens + 1 + ] + self.conv_state_shape, + dtype=self.conv_state_dtype, + device="cuda", + ) + self.intermediate_state_indices = torch.arange(max_batch_size, + dtype=torch.int32, + device="cuda") + + def _setup_replay_buffers(self, spec_config) -> None: + self.prev_num_accepted_tokens = None + self.cache_buf_idx = None + self.mamba_ssm_rand_seed = None + self.old_x = None + self.old_B = None + self.old_dt = None + self.old_dA_cumsum = None + + if (self.local_num_mamba_layers == 0 + or (not self._use_replay_state_update + and not self._mamba_ssm_stochastic_rounding)): + return + + cache_size = self.all_ssm_states[0].shape[0] + assert all(t.shape[0] == cache_size for t in self.all_ssm_states) + device = self.all_ssm_states[0].device + self.mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( + cache_size, self._seed_rank_offset, device) + + if spec_config is None or not self._use_replay_state_update: + return + + T = spec_config.tokens_per_gen_step + nheads, head_dim, d_state = self.ssm_state_shape + self.prev_num_accepted_tokens = torch.zeros(cache_size, + dtype=torch.int32, + device=device) + self.cache_buf_idx = torch.zeros(cache_size, + dtype=torch.int32, + device=device) + self.old_x = torch.zeros(self.local_num_mamba_layers, + cache_size, + 2, + T, + nheads, + head_dim, + dtype=self.conv_state_dtype, + device=device) + self.old_B = torch.zeros(self.local_num_mamba_layers, + cache_size, + 2, + T, + self._n_groups_per_rank, + d_state, + dtype=self.conv_state_dtype, + device=device) + self.old_dt = torch.zeros(self.local_num_mamba_layers, + cache_size, + 2, + nheads, + T, + dtype=torch.float32, + device=device) + self.old_dA_cumsum = torch.zeros(self.local_num_mamba_layers, + cache_size, + 2, + nheads, + T, + dtype=torch.float32, + device=device) + + def _attention_cache_bytes_per_token(self) -> int: + # Mamba layers have zero KV heads, so the generic calculation naturally + # returns only bytes owned by local attention layers. + return super().get_cache_bytes_per_token() + + def get_cache_bytes_per_token(self) -> int: + cache_bytes = self._attention_cache_bytes_per_token() + + interval = ( + self.kv_cache_config.mamba_state_config.periodic_snapshot_interval) + if (self.kv_cache_config.enable_block_reuse and interval is not None + and interval > 0): + cache_bytes += (self.local_num_mamba_layers * + (self.ssm_bytes + self.conv_bytes) // interval) + if cache_bytes == 0 and self.local_num_mamba_layers > 0: + return max( + 1, + self.local_num_mamba_layers * + (self.ssm_bytes + self.conv_bytes)) + return max(1, cache_bytes) + + def get_num_free_blocks(self) -> int: + assert len(self.kv_cache_map) == 0, ( + "get_num_free_blocks is only used when the kv cache manager is empty" + ) + attention_pages = [] + ssm_pages = [] + for local_layer_idx in range(self.num_local_layers): + layer_id = LayerId(local_layer_idx) + if self._is_local_mamba_layer(local_layer_idx): + ssm_pages.append( + self.impl.get_page_index_upper_bound( + layer_id, MambaRole.SSM_STATE) // + self._ssm_page_index_scale) + else: + attention_pages.append( + self.impl.get_page_index_upper_bound(layer_id, Role.KEY) // + self.kv_factor) + if attention_pages: + return max(attention_pages) + return max(ssm_pages) if ssm_pages else 0 + + @property + def blocks_in_primary_pool(self) -> int: + for local_layer_idx in range(self.num_local_layers): + if self._is_local_mamba_layer(local_layer_idx): + continue + return self.impl.get_page_index_upper_bound( + LayerId(local_layer_idx), Role.KEY) + return 0 + + def get_buffers(self, + layer_idx: int, + kv_layout: str = "NHD") -> Optional[torch.Tensor]: + local_layer_idx = self.layer_offsets[layer_idx] + if self._is_local_mamba_layer(local_layer_idx): + return None + return super().get_buffers(layer_idx, kv_layout) + + def check_invalid_values_in_kv_cache(self, + fill_with_zero: bool = False) -> bool: + """Check both attention pages and recurrent-state storage.""" + some_checks_unavailable = False + has_invalid_values = torch.tensor([False], + dtype=torch.bool, + device=torch.cuda.current_device()) + + def check_buffer(buffer: torch.Tensor) -> None: + nonlocal some_checks_unavailable + for start in range(0, buffer.shape[0], 256): + buffer_slice = buffer[start:start + 256] + try: + has_invalid_values.logical_or_( + torch.isnan(buffer_slice).any()) + has_invalid_values.logical_or_( + torch.isinf(buffer_slice).any()) + except NotImplementedError: + some_checks_unavailable = True + if fill_with_zero: + buffer.zero_() + + for global_layer_id, local_layer_id in self.layer_offsets.items(): + if self._is_local_mamba_layer(local_layer_id): + continue + # A layer group is a lifecycle, not a physical memory pool. + # Differently sized attention buffers can share one lifecycle, + # so scan every attention layer in this diagnostic path. + check_buffer(KVCacheManagerV2.get_buffers(self, global_layer_id)) + + for state_buffer in self.all_ssm_states: + check_buffer(state_buffer) + for state_buffer in self.all_conv_states: + check_buffer(state_buffer) + + torch.cuda.synchronize() + if some_checks_unavailable: + logger.warning( + "`torch.isnan` or `torch.isinf` is not implemented for a " + "configured cache dtype; related checks were skipped") + return bool(has_invalid_values) + + def add_dummy_requests( + self, + request_ids: List[int], + token_nums: Optional[List[int]] = None, + is_gen: bool = False, + prepare_resource: bool = True, + max_num_draft_tokens: int = 0, + kv_reserve_draft_tokens: Optional[int] = None, + use_mrope: bool = False, + max_beam_width: int = 1, + encoder_output_lens: Optional[List[int]] = None, + num_extra_decoding_steps: int = 0, + draft_kv_cache_manager: Optional[BaseResourceManager] = None, + ) -> List[LlmRequest]: + requests = super().add_dummy_requests( + request_ids=request_ids, + token_nums=token_nums, + is_gen=is_gen, + prepare_resource=prepare_resource, + max_num_draft_tokens=max_num_draft_tokens, + kv_reserve_draft_tokens=kv_reserve_draft_tokens, + use_mrope=use_mrope, + max_beam_width=max_beam_width, + encoder_output_lens=encoder_output_lens, + num_extra_decoding_steps=num_extra_decoding_steps, + draft_kv_cache_manager=draft_kv_cache_manager, + ) + if requests: + self.requests.extend(requests) + if prepare_resource: + # self.requests may still contain transfer-pending requests + # from an older batch. Only the new dummy requests participate + # in this forward pass. + self._setup_state_indices(requests) + return requests + + def free_resources(self, request: LlmRequest, pin_on_release: bool = False): + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is not None and kv_cache.is_active: + self.try_commit_blocks(request, kv_cache) + if request in self.requests: + self.requests.remove(request) + self._request_id_to_state_index.pop(request.py_request_id, None) + super().free_resources(request, pin_on_release) + + def prepare_resources(self, scheduled_batch: ScheduledRequests): + super().prepare_resources(scheduled_batch) + if self.local_num_mamba_layers == 0: + return + self.requests = (scheduled_batch.context_requests + + scheduled_batch.generation_requests) + self._setup_state_indices() + num_contexts = len(scheduled_batch.context_requests) + if num_contexts == 0: + return + ctx_slots = self.cuda_state_indices[:num_contexts].long() + if self._use_replay_state_update and self.prev_num_accepted_tokens is not None: + self.prev_num_accepted_tokens[ctx_slots] = 0 + self.cache_buf_idx[ctx_slots] = 0 + if self.old_x is not None: + self.old_x[:, ctx_slots] = 0 + if self.old_B is not None: + self.old_B[:, ctx_slots] = 0 + if self.old_dt is not None: + self.old_dt[:, ctx_slots] = 0 + if self.old_dA_cumsum is not None: + self.old_dA_cumsum[:, ctx_slots] = 0 + if self.mamba_ssm_rand_seed is not None: + self._seed_request_counter += 1 + counter = self._seed_request_counter + rank_offset = self._seed_rank_offset + host_slots = ctx_slots.cpu().tolist() + new_seeds = [ + _compute_deterministic_mamba_seed(counter, slot, rank_offset) + for slot in host_slots + ] + seed_tensor = torch.tensor( + new_seeds, + dtype=torch.int64, + device=self.mamba_ssm_rand_seed.device, + ) + self.mamba_ssm_rand_seed[ctx_slots] = seed_tensor + + def _setup_state_indices(self, + requests: Optional[List[LlmRequest]] = None + ) -> None: + if self.local_num_mamba_layers == 0: + return + requests = self.requests if requests is None else requests + n = len(requests) + assert n <= self._host_state_indices.shape[0], ( + f"State-index batch size {n} exceeds max_batch_size " + f"{self._host_state_indices.shape[0]}") + self._host_state_indices.zero_() + if n > 0: + for i, req in enumerate(requests): + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is None: + raise RuntimeError( + f"Missing V2 KV cache for request {req.py_request_id}") + base_index = kv_cache.get_ssm_block_base_index( + self.ssm_layer_group_id) + if base_index < 0: + raise RuntimeError( + f"Invalid SSM state block index {base_index} for " + f"request {req.py_request_id}") + self._host_state_indices[i] = base_index + + self.cuda_state_indices.copy_(self._host_state_indices, + non_blocking=True) + for i, req in enumerate(requests): + self._request_id_to_state_index[ + req.py_request_id] = self._host_state_indices[i].item() + + def get_state_indices(self, + request_ids: Optional[List[int]] = None, + is_padding: Optional[List[bool]] = None): + if self.local_num_mamba_layers == 0: + # Mamba metadata is still prepared on attention-only PP ranks, + # but no local kernel consumes these indices. Return harmless + # placeholders instead of consulting an intentionally empty map. + if request_ids is not None: + return [0] * len(request_ids) + return self.cuda_state_indices + if request_ids is not None: + return [self._request_id_to_state_index[rid] for rid in request_ids] + return self.cuda_state_indices + + def get_max_resource_count(self) -> int: + return self.max_batch_size + + def is_speculative(self) -> bool: + return self.spec_config is not None + + def update_mamba_states(self, + attn_metadata: "AttentionMetadata", + num_accepted_tokens: torch.Tensor, + state_indices: Optional[torch.Tensor] = None): + if self.local_num_mamba_layers == 0: + return + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + num_accepted_draft_tokens = ( + num_accepted_tokens[num_contexts:num_contexts + num_gens] - 1).to( + torch.int32) + if state_indices is None: + state_indices = self.get_state_indices() + state_indices_d = state_indices[num_contexts:num_contexts + + num_gens].to(torch.int32) + src_state_indices = self.intermediate_state_indices[:num_gens] + + if self._use_replay_state_update: + # Mirror Mamba2Metadata's replay checkpoint predicate: only a + # checkpoint write resets PNAT and flips the active buffer. + slots = state_indices_d.long() + accepted = num_accepted_tokens[num_contexts:num_contexts + + num_gens].to( + self.prev_num_accepted_tokens. + dtype) + prev_num_accepted_tokens = self.prev_num_accepted_tokens[slots] + wrote_checkpoint = (prev_num_accepted_tokens + + self.replay_step_width + > self.replay_history_size) + self.prev_num_accepted_tokens[slots] = torch.where( + wrote_checkpoint, accepted, prev_num_accepted_tokens + accepted) + cache_buf_idx = self.cache_buf_idx[slots] + self.cache_buf_idx[slots] = torch.where(wrote_checkpoint, + 1 - cache_buf_idx, + cache_buf_idx) + else: + for layer_offset, dst in enumerate(self.all_ssm_states): + _promote_mamba_state_triton( + dst.unsqueeze(0), + self.intermediate_ssm_states[layer_offset:layer_offset + 1], + src_state_indices, + num_accepted_draft_tokens, + state_indices_d, + ) + + for layer_offset, dst in enumerate(self.all_conv_states): + _promote_mamba_state_triton( + dst.unsqueeze(0), + self.intermediate_conv_states[layer_offset:layer_offset + 1], + src_state_indices, + num_accepted_draft_tokens, + state_indices_d, + ) + + def get_ssm_states(self, layer_idx: int) -> torch.Tensor: + return self.all_ssm_states[self.mamba_layer_offsets[layer_idx]] + + def get_conv_states(self, layer_idx: int) -> torch.Tensor: + return self.all_conv_states[self.mamba_layer_offsets[layer_idx]] + + def get_intermediate_ssm_states(self, + layer_idx: int) -> Optional[torch.Tensor]: + if self.intermediate_ssm_states is None: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.intermediate_ssm_states[layer_offset] + + def get_intermediate_conv_states(self, + layer_idx: int) -> Optional[torch.Tensor]: + if self.intermediate_conv_states is None: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.intermediate_conv_states[layer_offset] + + def mamba_layer_cache( + self, layer_idx: int + ) -> Union[PythonMambaCacheManager.State, + PythonMambaCacheManager.SpeculativeState, None]: + conv = self.get_conv_states(layer_idx) + ssm = self.get_ssm_states(layer_idx) + if self.spec_config is not None: + layer_offset = self.mamba_layer_offsets[layer_idx] + spec_kwargs = {} + if self.mamba_ssm_rand_seed is not None: + spec_kwargs['mamba_ssm_rand_seed'] = self.mamba_ssm_rand_seed + if self._use_replay_state_update: + spec_kwargs['old_x'] = self.old_x[layer_offset] + spec_kwargs['old_B'] = self.old_B[layer_offset] + spec_kwargs['old_dt'] = self.old_dt[layer_offset] + spec_kwargs['old_dA_cumsum'] = self.old_dA_cumsum[layer_offset] + spec_kwargs['cache_buf_idx'] = self.cache_buf_idx + spec_kwargs['prev_num_accepted_tokens'] = ( + self.prev_num_accepted_tokens) + else: + spec_kwargs['intermediate_ssm'] = self.intermediate_ssm_states[ + layer_offset] + return PythonMambaCacheManager.SpeculativeState( + conv=conv, + temporal=ssm, + intermediate_conv_window=self. + intermediate_conv_states[layer_offset], + **spec_kwargs, + ) + return PythonMambaCacheManager.State(conv=conv, temporal=ssm) + + def prepare_expect_snapshot_points(self, + requests: List[LlmRequest]) -> None: + """Set reusable Mamba snapshot boundaries before scheduling.""" + if not self.kv_cache_config.enable_block_reuse: + for request in requests: + request.expect_snapshot_points = [] + return + + state_config = self.kv_cache_config.mamba_state_config + interval = state_config.periodic_snapshot_interval + + for request in requests: + snapshot_points = set() + if interval is not None and interval > 0: + snapshot_points.update( + range(interval, request.prompt_len + 1, interval)) + for offset in state_config.additional_snapshot_offsets_from_start: + if offset <= request.prompt_len: + snapshot_points.add(offset) + for offset in state_config.additional_snapshot_offsets_from_end: + point = request.prompt_len - offset + if point > 0: + snapshot_points.add(point) + request.expect_snapshot_points = sorted(snapshot_points) + + def _mark_context_position_as_history(self, request: LlmRequest, + kv_cache) -> None: + """Advance history without making later recurrent state reusable.""" + history_length = request.context_current_position + if history_length <= kv_cache.history_length: + return + capacity = max(kv_cache.capacity, history_length) + if not kv_cache.resize(capacity, history_length=history_length): + raise ValueError( + "Failed to resize history length of V2 Mamba cache for " + f"request {request.py_request_id} to {history_length} tokens") + + def try_commit_blocks(self, request: LlmRequest, kv_cache=None) -> None: + should_block_reuse = (self.enable_block_reuse and not self.is_draft + and not request.is_dummy_request) + if not should_block_reuse: + return + + if kv_cache is None: + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is None: + return + + snapshot_points = request.expect_snapshot_points + commit_limit = (min(max(snapshot_points), request.prompt_len) + if snapshot_points else request.prompt_len) + commit_end = min(request.context_current_position, commit_limit) + if (self._is_block_reuse_commit_boundary(request) + and commit_end > kv_cache.num_committed_tokens): + tokens = self._augment_tokens_for_block_reuse( + request.get_tokens(DEFAULT_BEAM_INDEX), + request, + start=kv_cache.num_committed_tokens, + end=commit_end, + ) + kv_cache.commit(tokens) + if request.context_current_position >= commit_limit: + self._mark_context_position_as_history(request, kv_cache) + if request.context_remaining_length == 0: + kv_cache.stop_committing() + + def update_context_resources(self, + scheduled_batch: ScheduledRequests) -> None: + for request in scheduled_batch.context_requests: + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is None or not kv_cache.is_active: + continue + + should_block_reuse = (self.enable_block_reuse and not self.is_draft + and not request.is_dummy_request) + is_all_reusable = ( + self.block_reuse_policy == BlockReusePolicy.ALL_REUSABLE) + is_snapshot_boundary = (request.context_current_position + in request.expect_snapshot_points) + has_pending_snapshot = any( + point > request.context_current_position + for point in request.expect_snapshot_points) + should_resize = (not should_block_reuse or + (not is_all_reusable and not has_pending_snapshot)) + should_commit = (is_all_reusable or is_snapshot_boundary + or request.context_remaining_length == 0) + + if should_resize and not kv_cache.resize( + None, request.context_current_position): + raise ValueError( + "Failed to resize history length of V2 Mamba cache for " + f"request {request.py_request_id} to " + f"{request.context_current_position} tokens at context " + "update") + if should_commit: + self.try_commit_blocks(request, kv_cache) + if request.context_remaining_length == 0: + kv_cache.enable_swa_scratch_reuse = False + + def _is_block_reuse_commit_boundary(self, request: LlmRequest) -> bool: + """Snapshot recurrent state only at an explicitly prepared point.""" + return request.context_current_position in request.expect_snapshot_points + + @property + def use_replay_state_update(self) -> bool: + return self.get_replay_state_update_metadata() is not None + + def get_replay_state_update_metadata( + self) -> Optional[ReplayStateUpdateMetadata]: + prev_num_accepted_tokens = getattr(self, 'prev_num_accepted_tokens', + None) + cache_buf_idx = getattr(self, 'cache_buf_idx', None) + if (not self._use_replay_state_update + or prev_num_accepted_tokens is None or cache_buf_idx is None + or self.replay_step_width is None + or self.replay_history_size is None): + return None + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=prev_num_accepted_tokens, + cache_buf_idx=cache_buf_idx, + replay_step_width=self.replay_step_width, + replay_history_size=self.replay_history_size) + + def get_mamba_ssm_cache_dtype(self) -> torch.dtype: + return self.ssm_state_dtype + + def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: + return getattr(self, 'mamba_ssm_rand_seed', None) + + def shutdown(self): + self.all_ssm_states = [] + self.all_conv_states = [] + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None + self.prev_num_accepted_tokens = None + self.cache_buf_idx = None + self.mamba_ssm_rand_seed = None + self.old_x = None + self.old_B = None + self.old_dt = None + self.old_dA_cumsum = None + super().shutdown() diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index f82207881c8f..08ceedf61054 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -770,8 +770,9 @@ def drafting_loop_wrapper(model): ctx_chunk_config = None if kv_cache_config.enable_block_reuse and is_hybrid_linear(config): - ctx_chunk_config = (ContextChunkingPolicy.FORCE_CHUNK, - kv_cache_config.mamba_state_cache_interval) + # Snapshot boundaries come from expect_snapshot_points. The unit is + # only used to align chunks shortened by the scheduling budget. + ctx_chunk_config = (ContextChunkingPolicy.FORCE_CHUNK, tokens_per_block) guided_decoder: Optional[GuidedDecoder] = None if guided_decoding_config is not None: diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 157e2d7d720a..037b5a93e086 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -464,10 +464,13 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], pp_size = self.mapping.pp_size if self.mapping is not None else 1 live_state_slots = self.max_batch_size * pp_size max_snapshots = live_state_slots - if kv_cache_config.enable_block_reuse: - max_snapshots += ( - kv_cache_config.max_tokens // - linear_attention_metadata.states_snapshot_interval) + snapshot_interval = ( + linear_attention_metadata.states_snapshot_interval) + if (kv_cache_config.enable_block_reuse + and snapshot_interval is not None + and snapshot_interval > 0): + max_snapshots += (kv_cache_config.max_tokens // + snapshot_interval) blocks_per_window[LinearCacheType.RECURRENT_STATES.value] = ( int(max_snapshots), 0) @@ -1929,12 +1932,23 @@ def _calculate_max_num_blocks_for_linear_attention( pp_size = self.mapping.pp_size if self.mapping is not None else 1 intercept = self.max_batch_size * pp_size * state_bytes_local - max_tokens = max((primary_budget - intercept) // slope, 0) + if slope > 0: + max_tokens = max((primary_budget - intercept) // slope, 0) + elif primary_budget >= intercept: + # With snapshots disabled, a rank containing only recurrent-state + # layers has no per-token cache cost after its live slots are + # allocated. Bound the otherwise-unlimited token count by the + # configured capacity or by all resident sequences at max length. + max_tokens = (kv_cache_config.max_tokens + if kv_cache_config.max_tokens is not None else + self.max_seq_len * self.max_batch_size * pp_size) + else: + max_tokens = 0 if kv_cache_config.max_tokens is not None: max_tokens = min(kv_cache_config.max_tokens, max_tokens) if max_tokens < kv_cache_config.max_tokens: logger.warning( - f'The memory budget for Mamba + KV cache cannot fit the user-specified max_tokens of {kv_cache_config.max_tokens}. The calculated max_tokens based on the memory budget is {max_tokens}. Please consider adjusting max_batch_size/max_tokens/mamba_state_cache_interval.' + f'The memory budget for Mamba + KV cache cannot fit the user-specified max_tokens of {kv_cache_config.max_tokens}. The calculated max_tokens based on the memory budget is {max_tokens}. Please consider adjusting max_batch_size/max_tokens/mamba_state_config.periodic_snapshot_interval.' ) kv_blocks_in_primary_pool = int(max_tokens // self.tokens_per_block) diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index e4ec7ac25137..025cd8634d61 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -16,9 +16,9 @@ DynamicBatchConfig, Eagle3DecodingConfig, EagleDecodingConfig, EncodeCudaGraphConfig, ExtendedRuntimePerfKnobConfig, KvCacheConfig, LlmArgs, - LookaheadDecodingConfig, MedusaDecodingConfig, - MiniMaxM3SparseAttentionConfig, MoeConfig, - MTPDecodingConfig, NGramDecodingConfig, + LookaheadDecodingConfig, MambaStateConfig, + MedusaDecodingConfig, MiniMaxM3SparseAttentionConfig, + MoeConfig, MTPDecodingConfig, NGramDecodingConfig, PARDDecodingConfig, PrometheusMetricsConfig, ReorderRequestPolicyConfig, RocketSparseAttentionConfig, SADecodingConfig, SAEnhancerConfig, @@ -43,6 +43,7 @@ 'ConversationParams', 'DisaggScheduleStyle', 'KvCacheConfig', + 'MambaStateConfig', 'KvCacheRetentionConfig', 'CudaGraphConfig', 'DecodeCudaGraphConfig', diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 6cdd95827ccf..3e849cc8df12 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -34,7 +34,7 @@ from pydantic import AliasChoices, BaseModel, ConfigDict from pydantic import Field as PydanticField from pydantic import (NonNegativeFloat, NonNegativeInt, PositiveInt, - PrivateAttr, field_validator, model_validator) + PrivateAttr, StrictInt, field_validator, model_validator) from strenum import StrEnum from transformers import PreTrainedTokenizerBase @@ -58,9 +58,7 @@ CapacitySchedulerPolicy as _CapacitySchedulerPolicy, ContextChunkingPolicy as _ContextChunkingPolicy, DecodingConfig, - DecodingMode, DynamicBatchConfig as _DynamicBatchConfig, - EagleConfig as _EagleConfig, ExecutorConfig as _ExecutorConfig, ExtendedRuntimePerfKnobConfig as _ExtendedRuntimePerfKnobConfig, KvCacheConfig as _KvCacheConfig, @@ -3460,6 +3458,85 @@ class ReorderRequestPolicyConfig(StrictBaseModel): description="The arguments of the request reordering policy.") +_StrictPositiveInt = Annotated[StrictInt, PydanticField(gt=0)] +_StrictNonNegativeInt = Annotated[StrictInt, PydanticField(ge=0)] + + +class MambaStateConfig(StrictBaseModel): + """Configuration for reusable Mamba recurrent-state snapshots.""" + + periodic_snapshot_interval: NonNegativeInt = Field( + default=256, + description= + "The number of tokens between periodic snapshots in the Mamba " + "prefix cache. Set to 0 to disable periodic snapshots.") + + additional_snapshot_offsets_from_start: List[_StrictPositiveInt] = Field( + default_factory=list, + status="prototype", + telemetry=False, + description= + "Additional Mamba state snapshot offsets measured from the start " + "of each prompt. Offsets beyond the prompt length are ignored. " + "These snapshots require KV cache manager V2.") + + additional_snapshot_offsets_from_end: List[_StrictNonNegativeInt] = Field( + default_factory=list, + status="prototype", + telemetry=False, + description= + "Additional Mamba state snapshot offsets measured backward from " + "the end of each prompt. An offset of 0 selects the prompt end. " + "Offsets that do not resolve inside the prompt are ignored. These " + "snapshots require KV cache manager V2.") + + +def _migrate_legacy_mamba_interval_from_config_file( + config_dict: Dict[str, Any]) -> Dict[str, Any]: + """Migrate the legacy Mamba interval in a YAML/JSON config mapping. + + This intentionally runs only at config-file ingestion boundaries. The + Pydantic models remain strict so direct Python construction accepts only + ``mamba_state_config.periodic_snapshot_interval``. + """ + kv_cache_config = config_dict.get("kv_cache_config") + legacy_field = "mamba_state_cache_interval" + if not isinstance(kv_cache_config, + dict) or legacy_field not in kv_cache_config: + return config_dict + + migrated_config = dict(config_dict) + migrated_kv_cache_config = dict(kv_cache_config) + migrated_config["kv_cache_config"] = migrated_kv_cache_config + + state_config_field = "mamba_state_config" + interval_field = "periodic_snapshot_interval" + if state_config_field in migrated_kv_cache_config: + state_config = migrated_kv_cache_config[state_config_field] + if not isinstance(state_config, dict): + raise ValueError( + "Config file field 'kv_cache_config.mamba_state_config' must " + "be a mapping when using the legacy " + "'kv_cache_config.mamba_state_cache_interval' option.") + if interval_field in state_config: + raise ValueError("Config file cannot set both " + "'kv_cache_config.mamba_state_cache_interval' and " + "'kv_cache_config.mamba_state_config." + "periodic_snapshot_interval'.") + migrated_state_config = dict(state_config) + else: + migrated_state_config = {} + + migrated_state_config[interval_field] = migrated_kv_cache_config.pop( + legacy_field) + migrated_kv_cache_config[state_config_field] = migrated_state_config + logger.warning( + "Config-file option 'kv_cache_config.mamba_state_cache_interval' is " + "deprecated; use 'kv_cache_config.mamba_state_config." + "periodic_snapshot_interval' instead.") + return migrated_config + + @PybindMirror.mirror_pybind_fields(_KvCacheConfig) class KvCacheConfig(StrictBaseModel, PybindMirror): """Configuration for the KV cache.""" @@ -3593,10 +3670,9 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): description="The number of tokens per block.") # This is a pure python field, not a pybind field. It is only for the Pytorch backend. - mamba_state_cache_interval: PositiveInt = Field( - default=256, - description= - "The number of tokens between cache steps in the Mamba prefix cache.") + mamba_state_config: MambaStateConfig = Field( + default_factory=MambaStateConfig, + description="Configuration for reusable Mamba state snapshots.") use_kv_cache_manager_v2: bool | Literal["auto"] = Field( default="auto", @@ -4326,6 +4402,8 @@ def speculative_model(self) -> Optional[Union[str, Path]]: def from_yaml(cls, yaml_path: Union[str, Path]): with open(yaml_path, "r") as f: config_dict = yaml.safe_load(f) + config_dict = _migrate_legacy_mamba_interval_from_config_file( + config_dict) return cls(**config_dict) @field_validator("dtype") @@ -5753,6 +5831,9 @@ def update_llm_args_with_extra_dict( If `explicit_cli_keys` is None, YAML wins on conflicts. """ + llm_args_dict = _migrate_legacy_mamba_interval_from_config_file( + llm_args_dict) + # CLI scalar -> nested KvCacheConfig field. Callers add the CLI scalar # name to `explicit_cli_keys` to make it win over YAML's same-named # field inside `kv_cache_config:`. diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 0945db4d2394..acdcfeb1d608 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -156,6 +156,7 @@ LayerConfig = AttentionLayerConfig | SsmLayerConfig class KVCacheDesc: capacity: int history_length: int + num_ssm_slots: int = 1 @dataclass(slots=True) class BatchDesc: diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py index 38798e17b03a..4d55539b6b05 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py @@ -361,6 +361,10 @@ def __init__(self, tokens: Sequence[TokenIdExt], prev: "Block | RootBlock") -> N assert NDEBUG or (not b.is_full and b is not self and b.key == k and not b.next) to_remove.append(k) event_manager = get_tree(prev).event_manager if to_remove else None + # Keep RootBlock attached while covered children are replaced. Adding + # the replacement first prevents detach_next() from pruning an + # otherwise-empty root before this block becomes its new child. + prev.next[self.key] = self for k in to_remove: b = detach_next(prev, k) assert isinstance(b, Block) @@ -368,7 +372,6 @@ def __init__(self, tokens: Sequence[TokenIdExt], prev: "Block | RootBlock") -> N event_manager.add_removed_event(b.key) assert b.is_orphan # _KVCache may still hold it. # prev.next keeps a strong ref to this _Block, so no need to remove self from prev.next in __del__(). - prev.next[self.key] = self def _release_pages(self) -> None: """Reclaim every page held by this block. diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py index 8f374e7a37f7..ed3912070985 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py @@ -145,9 +145,11 @@ def __post_init__(self) -> None: class KVCacheDesc: capacity: int history_length: int + num_ssm_slots: int = 1 def __post_init__(self) -> None: assert 0 <= self.history_length <= self.capacity + assert self.num_ssm_slots > 0 # A batch of requests, working as a use case the KVCacheManager must always support. diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py index 69836ab0efcf..6f1f1e9dfffc 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py @@ -956,12 +956,22 @@ def _compute_slots_for_batch( """Compute the minimum number of slots per pool group to support a BatchDesc.""" num_slots = filled_list(0, self.num_pool_groups) ssm_lc_idx = self._life_cycles.ssm_life_cycle_id + total_ssm_slots = sum(kv.num_ssm_slots for kv in batch.kv_caches) + num_request_lineages = len(batch.kv_caches) + # A non-live SSM slot indicates that this batch can retain snapshots. + # Snapshot alignment is unknown while storage is being sized, so once + # such capacity exists, reserve one partial attention page per request + # lineage. V2 replaces covered partial pages within a lineage, making N + # a safe upper bound. This does not require S >= 2N: when N < S < 2N, + # the bound merely over-reserves attention pages. The guard keeps the + # default num_ssm_slots=1 descriptors used by non-Mamba managers from + # paying this Mamba-specific overhead. + has_non_live_ssm_capacity = total_ssm_slots > num_request_lineages sys_blocks = batch.system_prompt_length // tokens_per_block for lc_idx, lc in typed_enumerate(self._life_cycles.get()): pg_idx = self.get_pool_group_index(lc_idx) if lc_idx == ssm_lc_idx: - # SSM: always 1 dedicated block per request, never shared. - num_slots[pg_idx] += len(batch.kv_caches) + num_slots[pg_idx] += total_ssm_slots continue # Shared sys blocks (counted once): union of non-stale sys blocks across all requests. # A sys block needs memory if it's non-stale for ANY request. @@ -998,6 +1008,11 @@ def _compute_slots_for_batch( ) else: num_slots[pg_idx] += unique_non_stale + # Retained partial-page copies are additional physical blocks, not + # tokens in KVCacheDesc.capacity, so add their safe upper bound + # after token staleness and scratch-sharing calculations. + if has_non_live_ssm_capacity: + num_slots[pg_idx] += num_request_lineages return num_slots def _slots_to_bytes( diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 739db09bfb23..6eefc84416df 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -710,7 +710,7 @@ "annotation": "", "converter": "", "kind": "value", - "path": "kv_cache_config.mamba_state_cache_interval" + "path": "kv_cache_config.mamba_state_config.periodic_snapshot_interval" }, { "allowed_values": [], diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 0bf4298efaad..2c62c5bb2492 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -31,11 +31,11 @@ from tensorrt_llm.llmapi import ( AttentionDpConfig, CudaGraphConfig, DeepSeekSparseAttentionConfig, DFlashDecodingConfig, DSparkDecodingConfig, DraftTargetDecodingConfig, - Eagle3DecodingConfig, KvCacheConfig, MiniMaxM3SparseAttentionConfig, - MoeConfig, MTPDecodingConfig, NGramDecodingConfig, PARDDecodingConfig, - RocketSparseAttentionConfig, SADecodingConfig, SamplingParams, - SchedulerConfig, SkipSoftmaxAttentionConfig, SAEnhancerConfig, - TorchCompileConfig) + Eagle3DecodingConfig, KvCacheConfig, MambaStateConfig, + MiniMaxM3SparseAttentionConfig, MoeConfig, MTPDecodingConfig, + NGramDecodingConfig, PARDDecodingConfig, RocketSparseAttentionConfig, + SADecodingConfig, SamplingParams, SchedulerConfig, + SkipSoftmaxAttentionConfig, SAEnhancerConfig, TorchCompileConfig) # isort: on from tensorrt_llm.quantization import QuantAlgo @@ -5899,7 +5899,7 @@ def test_nvfp4(self, moe_backend, tp_size, ep_size, attention_dp, kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.6, enable_block_reuse=enable_block_reuse) if enable_block_reuse: - kv_cache_config.mamba_state_cache_interval = 256 + kv_cache_config.mamba_state_config.periodic_snapshot_interval = 256 pytorch_config = dict(disable_overlap_scheduler=False, cuda_graph_config=CudaGraphConfig( max_batch_size=512, enable_padding=True)) @@ -7062,7 +7062,7 @@ def test_fp8_4gpus(self, attention_dp, use_cpp_mamba, monkeypatch): @skip_pre_blackwell @pytest.mark.parametrize( - "tp_size, ep_size, mamba_state_cache_interval, attention_dp, use_mtp", + "tp_size, ep_size, periodic_snapshot_interval, attention_dp, use_mtp", [ (1, 1, 256, False, False), (4, 1, 256, False, True), @@ -7073,7 +7073,7 @@ def test_fp8_4gpus(self, attention_dp, use_cpp_mamba, monkeypatch): ids=["TP1", "TP4_MTP", "TEP4", "DEP4_MTP_OFF", "DEP4_MTP_ON"], ) def test_nvfp4_4gpus_block_reuse(self, tp_size, ep_size, - mamba_state_cache_interval, attention_dp, + periodic_snapshot_interval, attention_dp, use_mtp): gpu_needed = max(tp_size, ep_size) if get_device_count() < gpu_needed: @@ -7089,7 +7089,8 @@ def test_nvfp4_4gpus_block_reuse(self, tp_size, ep_size, kv_cache_config=KvCacheConfig( enable_block_reuse=True, mamba_ssm_cache_dtype="float16", - mamba_state_cache_interval=mamba_state_cache_interval, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=periodic_snapshot_interval), free_gpu_memory_fraction=0.8, ), max_batch_size=32, @@ -7495,7 +7496,7 @@ def test_nvfp4_8gpus(self, attention_dp, moe_backend): @skip_pre_blackwell @pytest.mark.skip_less_mpi_world_size(4) @pytest.mark.parametrize( - "tp_size, ep_size, mamba_state_cache_interval, attention_dp, use_mtp", + "tp_size, ep_size, periodic_snapshot_interval, attention_dp, use_mtp", [ (4, 1, 256, False, True), (4, 4, 256, False, False), @@ -7505,7 +7506,7 @@ def test_nvfp4_8gpus(self, attention_dp, moe_backend): ids=["TP4_MTP", "TEP4", "ADP4", "ADP4_MTP"], ) def test_nvfp4_4gpus_block_reuse(self, tp_size, ep_size, - mamba_state_cache_interval, attention_dp, + periodic_snapshot_interval, attention_dp, use_mtp): mtp_config = MTPDecodingConfig( num_nextn_predict_layers=3, @@ -7517,7 +7518,8 @@ def test_nvfp4_4gpus_block_reuse(self, tp_size, ep_size, kv_cache_config=KvCacheConfig( enable_block_reuse=True, mamba_ssm_cache_dtype="float16", - mamba_state_cache_interval=mamba_state_cache_interval, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=periodic_snapshot_interval), free_gpu_memory_fraction=0.5, ), max_batch_size=max_batch_size, diff --git a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py index 307588e0d2a0..59c3c3515f6e 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py @@ -408,3 +408,21 @@ def test_per_conversation_policy_ignores_overlapping_request( _free_if_active(manager, request_old_prompt) _free_if_active(manager, request_b) _free_if_active(manager, request_a) + + +def test_iteration_stats_reports_physical_pool_groups_without_window_metadata() -> None: + manager = object.__new__(KVCacheManagerV2) + manager.enable_stats = True + manager.impl = SimpleNamespace( + cache_tier_list=[object()], + get_and_reset_iteration_stats=lambda: {}, + ) + manager._stats_life_cycle_metadata = lambda: {} + manager._storage_pool_groups_by_window = lambda: {} + manager._get_and_reset_iteration_peak_block_stats = lambda _level: [None, None] + manager._get_storage_statistics = lambda _level: [object(), object()] + manager._build_pool_group_iteration_stats = lambda pool_group_id, *_args: pool_group_id + + stats = manager.get_iteration_stats() + + assert stats.by_pool_group == {0: 0, 1: 1} diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py index 1e545ac5cb79..0a7c67e9e702 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py @@ -1542,6 +1542,21 @@ def test_force_chunk_allows_snapshot_point_shorter_than_unit_size(self): assert ids(out.context_requests) == [0] assert req.context_chunk_size == 150 + def test_force_chunk_makes_progress_before_budget_limited_snapshot(self): + mgr = make_kv_cache_manager(tokens_per_block=32) + sched = make_scheduler( + mgr, + max_num_tokens=128, + ctx_chunk_config=(ContextChunkingPolicy.FORCE_CHUNK, 32), + ) + req = make_ctx_request(0, context_remaining_length=512) + req.expect_snapshot_points = [256, 512] + + out = sched.schedule_request([req], set()) + + assert ids(out.context_requests) == [0] + assert req.context_chunk_size == 128 + def test_multiple_ctx_share_budget(self): """Two ctx requests share the budget.""" mgr = make_kv_cache_manager(tokens_per_block=64) diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 0291ef3e74dd..7595f5b50dc7 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -1,15 +1,22 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Regression tests for MambaCacheManager padding-slot behavior and -CppMambaHybridCacheManager PP-sharding edge cases.""" +"""Regression tests for Python, Cpp, and V2 Mamba cache managers.""" import os from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import torch +from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata +from tensorrt_llm._torch.pyexecutor._util import ( + KvCacheCreator, + _create_kv_cache_manager, + get_kv_cache_manager_cls, +) from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import BlockReusePolicy, KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import ( ATTENTION_DP_DUMMY_REQUEST_ID, LlmRequest, @@ -18,19 +25,392 @@ from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( MIN_REPLAY_HISTORY_SIZE, CppMambaHybridCacheManager, + MambaRole, + MixedMambaHybridCacheManager, PythonMambaCacheManager, + V2MambaHybridCacheManager, + _get_local_mamba_cache_layout, _get_mamba_hybrid_pool_size, + _get_num_cuda_graph_padding_dummy_slots, + _mamba_snapshot_rule_counts, +) +from tensorrt_llm._torch.pyexecutor.resource_manager import ( + CacheTypeCpp, + DataType, + KVCacheManager, + get_pp_layers, ) -from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, DataType, KVCacheManager from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._utils import torch_dtype_to_binding from tensorrt_llm.bindings.internal.batch_manager import LinearCacheType -from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MTPDecodingConfig +from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MambaStateConfig, MTPDecodingConfig from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import GpuCacheTierConfig, LayerId +from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManager as RuntimeKVCacheManager skip_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_cuda_graph_padding_dummy_slot_count_tracks_reachable_draft_lengths(): + assert _get_num_cuda_graph_padding_dummy_slots(None, 64) == 1 + + static = MTPDecodingConfig(max_draft_len=4) + assert _get_num_cuda_graph_padding_dummy_slots(static, 64) == 1 + + gated = MTPDecodingConfig( + max_draft_len=4, + acceptance_rate_window_size=8, + acceptance_rate_threshold=0.5, + ) + assert _get_num_cuda_graph_padding_dummy_slots(gated, 64) == 2 + + dynamic = MTPDecodingConfig( + max_draft_len=4, + draft_len_schedule={4: 4, 8: 2, 32: 1}, + ) + assert _get_num_cuda_graph_padding_dummy_slots(dynamic, 5) == 2 + assert _get_num_cuda_graph_padding_dummy_slots(dynamic, 32) == 3 + assert _get_num_cuda_graph_padding_dummy_slots(dynamic, 64) == 4 + + repeated = MTPDecodingConfig( + max_draft_len=4, + draft_len_schedule={4: 4, 8: 4, 32: 2}, + ) + assert _get_num_cuda_graph_padding_dummy_slots(repeated, 64) == 3 + + +def _hybrid_model_config(): + config = SimpleNamespace( + architectures=["Qwen3_5MoeForCausalLM"], + num_hidden_layers=2, + layer_types=["linear_attention", "full_attention"], + ) + return SimpleNamespace( + pretrained_config=config, + sparse_attention_config=None, + get_num_mamba_layers=lambda: 1, + ) + + +def _hybrid_cache_sizing_model_config(layer_types): + config = SimpleNamespace( + architectures=["Qwen3_5ForCausalLM"], + num_hidden_layers=len(layer_types), + layer_types=layer_types, + linear_key_head_dim=8, + linear_conv_kernel_dim=4, + linear_num_value_heads=4, + linear_num_key_heads=1, + linear_value_head_dim=8, + num_key_value_heads=2, + num_attention_heads=4, + hidden_size=32, + torch_dtype=torch.float16, + ) + return SimpleNamespace(pretrained_config=config, quant_config=None) + + +@pytest.mark.parametrize( + ("use_v2", "enable_block_reuse", "expected"), + [ + (True, False, V2MambaHybridCacheManager), + (True, True, V2MambaHybridCacheManager), + (False, False, CppMambaHybridCacheManager), + (False, True, CppMambaHybridCacheManager), + ("auto", False, CppMambaHybridCacheManager), + ("auto", True, CppMambaHybridCacheManager), + ], +) +def test_hybrid_cache_manager_factory_honors_v2_setting( + monkeypatch, use_v2, enable_block_reuse, expected +): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + kv_cache_config = KvCacheConfig( + enable_block_reuse=enable_block_reuse, + use_kv_cache_manager_v2=use_v2, + ) + + assert get_kv_cache_manager_cls(_hybrid_model_config(), kv_cache_config) is expected + + +def test_qwen3_gdn_replay_falls_back_for_v2_manager(monkeypatch): + """GDN's all-layer replay commit requires the contiguous C++ state view.""" + captured_cpp = {} + captured_v2 = {} + + class RecordingCppManager(CppMambaHybridCacheManager): + def __init__(self, *args, **kwargs): + captured_cpp.update(kwargs) + + class RecordingV2Manager(V2MambaHybridCacheManager): + def __init__(self, *args, **kwargs): + captured_v2.update(kwargs) + + pretrained_config = SimpleNamespace( + architectures=["Qwen3_5MoeForCausalLM"], + hidden_size=32, + num_attention_heads=4, + num_key_value_heads=2, + num_hidden_layers=2, + layer_types=["linear_attention", "full_attention"], + ) + model_config = SimpleNamespace( + pretrained_config=pretrained_config, + quant_config=None, + ) + mamba_params = SimpleNamespace( + state_size=8, + conv_kernel=4, + num_heads=4, + n_groups=1, + head_dim=8, + num_mamba_layers=1, + mamba_layer_mask=[True, False], + dtype=torch.bfloat16, + mamba_ssm_cache_dtype=torch.bfloat16, + num_full_attention_layers=1, + full_attention_layer_mask=[False, True], + ) + monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "1") + monkeypatch.setattr("tensorrt_llm._torch.pyexecutor._util.get_sm_version", lambda: 90) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor._util.extract_mamba_kv_cache_params", + lambda *args, **kwargs: mamba_params, + ) + + common_kwargs = dict( + model_engine=None, + mapping=Mapping(world_size=1, tp_size=1, pp_size=1), + tokens_per_block=32, + max_seq_len=2048, + # Cover the batch size at which upstream selects partitioned replay. + max_batch_size=16, + spec_config=MTPDecodingConfig(max_draft_len=3), + sparse_attention_config=None, + max_num_tokens=256, + max_beam_width=1, + kv_connector_manager=None, + model_config=model_config, + dtype=torch.bfloat16, + is_draft=False, + ) + _create_kv_cache_manager( + kv_cache_manager_cls=RecordingCppManager, + kv_cache_config=KvCacheConfig(use_kv_cache_manager_v2=False), + **common_kwargs, + ) + _create_kv_cache_manager( + kv_cache_manager_cls=RecordingV2Manager, + kv_cache_config=KvCacheConfig(use_kv_cache_manager_v2=True), + **common_kwargs, + ) + + assert captured_cpp["use_replay_state_update"] is True + assert captured_cpp["model_type"] == "qwen3_next" + assert captured_v2["use_replay_state_update"] is False + assert "model_type" not in captured_v2 + + +def test_hybrid_cache_manager_factory_rejects_cpp_preference_with_explicit_v2( + monkeypatch, +): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.setenv("TLLM_MAMBA_MANAGER_PREFERENCE", "CPP") + + with pytest.raises(ValueError, match="conflicts with explicit"): + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + enable_block_reuse=True, + use_kv_cache_manager_v2=True, + ), + ) + + +def test_hybrid_cache_manager_factory_v2_preference_does_not_select_v2( + monkeypatch, +): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.setenv("TLLM_MAMBA_MANAGER_PREFERENCE", "V2") + + assert ( + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig(use_kv_cache_manager_v2=False), + ) + is CppMambaHybridCacheManager + ) + + +def test_hybrid_cache_manager_factory_requires_v2_for_explicit_snapshots( + monkeypatch, +): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + + with pytest.raises(ValueError, match="use_kv_cache_manager_v2=True"): + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + enable_block_reuse=False, + mamba_state_config=MambaStateConfig(additional_snapshot_offsets_from_end=[0]), + use_kv_cache_manager_v2=False, + ), + ) + + +def test_hybrid_cache_manager_factory_requires_snapshot_policy(monkeypatch): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + + with pytest.raises(ValueError, match="at least one snapshot policy"): + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=0), + use_kv_cache_manager_v2=True, + ), + ) + + +def test_hybrid_cache_manager_factory_routes_explicit_snapshots_to_v2( + monkeypatch, +): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + + assert ( + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=0, + additional_snapshot_offsets_from_end=[0], + ), + use_kv_cache_manager_v2=True, + ), + ) + is V2MambaHybridCacheManager + ) + + +@pytest.mark.parametrize( + ("env_name", "env_value", "expected_error"), + [ + ( + "TLLM_MAMBA_MANAGER_PREFERENCE", + "MIXED", + "does not support block reuse", + ), + ( + "TRTLLM_USE_PY_MAMBA", + "1", + "does not support block reuse", + ), + ], +) +def test_hybrid_cache_manager_factory_rejects_mixed_override_with_reuse( + monkeypatch, env_name, env_value, expected_error +): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + monkeypatch.setenv(env_name, env_value) + + with pytest.raises(ValueError, match=expected_error): + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig(enable_block_reuse=True), + ) + + +@pytest.mark.parametrize("use_v2", [False, "auto"]) +def test_hybrid_cache_manager_factory_keeps_v1_disagg_route(monkeypatch, use_v2): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + + assert ( + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + enable_block_reuse=False, + use_kv_cache_manager_v2=use_v2, + ), + is_disagg=True, + ) + is CppMambaHybridCacheManager + ) + assert ( + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + enable_block_reuse=False, + use_kv_cache_manager_v2=use_v2, + ), + is_disagg=True, + cache_transceiver_config=SimpleNamespace(transceiver_runtime="PYTHON"), + ) + is MixedMambaHybridCacheManager + ) + + +@pytest.mark.parametrize("enable_block_reuse", [False, True]) +@pytest.mark.parametrize("transceiver_runtime", [None, "PYTHON"]) +def test_hybrid_cache_manager_factory_rejects_v2_disagg( + monkeypatch, enable_block_reuse, transceiver_runtime +): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + transceiver_config = ( + None + if transceiver_runtime is None + else SimpleNamespace(transceiver_runtime=transceiver_runtime) + ) + + with pytest.raises(ValueError, match="V2.*not supported with disaggregated"): + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + enable_block_reuse=enable_block_reuse, + use_kv_cache_manager_v2=True, + ), + is_disagg=True, + cache_transceiver_config=transceiver_config, + ) + + +@pytest.mark.parametrize( + "max_beam_width, has_connector, expected", + [ + (2, False, "max_beam_width > 1"), + (1, True, "kv_connector_manager"), + (2, True, "kv_connector_manager, max_beam_width > 1"), + ], +) +def test_v2_hybrid_incompatibility_fails_without_cpp_fallback( + max_beam_width, has_connector, expected +): + config = SimpleNamespace( + architectures=["Qwen3_5MoeForCausalLM"], + num_hidden_layers=2, + layer_types=["linear_attention", "full_attention"], + ) + model_config = SimpleNamespace( + pretrained_config=config, + sparse_attention_config=None, + ) + creator = object.__new__(KvCacheCreator) + creator._kv_connector_manager = object() if has_connector else None + creator._max_beam_width = max_beam_width + + with pytest.raises(NotImplementedError, match=expected): + creator._fallback_if_unsupported_kv_cache_manager_v2( + V2MambaHybridCacheManager, model_config, KvCacheConfig() + ) + + def _make_mgr( max_batch_size=4, max_draft_len=2, enable_attention_dp=False, use_replay_state_update=False ): @@ -346,11 +726,420 @@ def test_non_mtp_pytorch_prepare_and_get_state_indices_flow(): ] +def test_v2_hybrid_prepare_expect_snapshot_points(): + mgr = object.__new__(V2MambaHybridCacheManager) + mgr.kv_cache_config = KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=64, + additional_snapshot_offsets_from_start=[128, 999], + additional_snapshot_offsets_from_end=[0, 22, 999], + ), + ) + requests = [ + SimpleNamespace(prompt_len=150, expect_snapshot_points=[999]), + SimpleNamespace(prompt_len=128, expect_snapshot_points=[]), + SimpleNamespace(prompt_len=32, expect_snapshot_points=[]), + ] + + mgr.prepare_expect_snapshot_points(requests) + + assert [request.expect_snapshot_points for request in requests] == [ + [64, 128, 150], + [64, 106, 128], + [10, 32], + ] + + +def test_v2_hybrid_prepare_expect_snapshot_points_without_periodic_snapshots(): + mgr = object.__new__(V2MambaHybridCacheManager) + mgr.kv_cache_config = KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=0, + additional_snapshot_offsets_from_start=[128], + additional_snapshot_offsets_from_end=[0, 13], + ), + ) + request = SimpleNamespace(prompt_len=150, expect_snapshot_points=[]) + + mgr.prepare_expect_snapshot_points([request]) + + assert request.expect_snapshot_points == [128, 137, 150] + + +def test_mamba_snapshot_rule_count_deduplicates_and_filters_unreachable_points(): + config = KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=0, + additional_snapshot_offsets_from_start=[64, 65, 64], + additional_snapshot_offsets_from_end=[0, 32, 4096], + ), + ) + + assert _mamba_snapshot_rule_counts(config, 128, 32) == (4, 3) + + +def test_v2_hybrid_snapshot_sizing_scales_with_pp_and_explicit_rules(): + mgr = object.__new__(V2MambaHybridCacheManager) + mgr.max_batch_size = 4 + mgr.mapping = Mapping(world_size=2, rank=0, tp_size=1, pp_size=2) + mgr.max_seq_len = 128 + mgr.tokens_per_block = 32 + mgr._num_reserved_dummy_slots = 0 + config = KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=0, + additional_snapshot_offsets_from_start=[64], + additional_snapshot_offsets_from_end=[0, 32], + ), + ) + + assert mgr._num_ssm_snapshots_for_capacity(512, config) == 24 + assert mgr._ssm_slots_per_request_for_typical_batch(512, config) == [4] * 8 + + periodic_config = KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=48), + ) + assert mgr._ssm_slots_per_request_for_typical_batch(47, periodic_config) == [1] * 8 + assert mgr._ssm_slots_per_request_for_typical_batch(48, periodic_config) == [2] + [1] * 7 + + +def test_cpp_mamba_estimator_handles_disabled_snapshots_without_attention(): + manager = object.__new__(KVCacheManager) + manager._primary_pool_memory_bytes = 4096 + manager._secondary_pool_memory_bytes = 0 + manager.linear_attention_metadata = SimpleNamespace( + all_recurrent_states_bytes=64, + states_snapshot_interval=0, + ) + manager.max_attention_window_vec = [LinearCacheType.RECURRENT_STATES.value] + manager.get_cache_bytes_per_token = lambda: 0 + manager.mapping = SimpleNamespace(pp_size=1) + manager.max_batch_size = 4 + manager.max_seq_len = 128 + manager.tokens_per_block = 32 + manager.spec_config = None + + blocks = manager._calculate_max_num_blocks_for_linear_attention( + KvCacheConfig( + max_tokens=512, + enable_block_reuse=False, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=0), + ) + ) + + assert blocks[128] == (16, 0) + assert blocks[LinearCacheType.RECURRENT_STATES.value] == (5, 0) + + +def test_hybrid_mtp_layout_honors_explicit_base_partition(): + model_config = _hybrid_cache_sizing_model_config( + [ + "linear_attention", + "full_attention", + "linear_attention", + "full_attention", + ] + ) + spec_config = MTPDecodingConfig(max_draft_len=1) + expected_pp_layers = ([0, 1], [2, 3, 4]) + + for rank in range(2): + mapping = Mapping( + world_size=2, + rank=rank, + tp_size=1, + pp_size=2, + pp_partition=[2, 2], + ) + pp_layers, total_layers = get_pp_layers( + 5, + mapping, + spec_config=spec_config, + layer_mask=[True] * 5, + ) + assert total_layers == 5 + assert pp_layers == expected_pp_layers[rank] + + _, local_mamba_layers, local_attention_layers = _get_local_mamba_cache_layout( + model_config, + mapping, + spec_config=spec_config, + ) + assert local_mamba_layers == 1 + assert local_attention_layers == rank + 1 + + for manager_cls in ( + CppMambaHybridCacheManager, + V2MambaHybridCacheManager, + ): + cache_cost = manager_cls.get_cache_size_per_token( + model_config, + mapping, + max_batch_size=1, + kv_cache_config=KvCacheConfig(enable_block_reuse=False), + spec_config=spec_config, + ) + extra_attention_bound = ( + 4096 * (rank + 1) if manager_cls is V2MambaHybridCacheManager else 0 + ) + assert cache_cost == (64 * (rank + 1), 2400 + extra_attention_bound) + + +def test_hybrid_separate_mtp_draft_estimator_has_no_mamba_state(): + model_config = _hybrid_cache_sizing_model_config( + [ + "linear_attention", + "full_attention", + "linear_attention", + "full_attention", + ] + ) + spec_config = MTPDecodingConfig(max_draft_len=1) + target_layer_mask = [True] * 4 + draft_layer_mask = [False] * 4 + [True] + + for rank in range(2): + mapping = Mapping(world_size=2, rank=rank, tp_size=1, pp_size=2) + + target_cost = V2MambaHybridCacheManager.get_cache_size_per_token( + model_config, + mapping, + max_batch_size=1, + kv_cache_config=KvCacheConfig(enable_block_reuse=False), + spec_config=spec_config, + layer_mask=target_layer_mask, + ) + assert target_cost == (64, 6496) + + _, local_mamba_layers, local_attention_layers = _get_local_mamba_cache_layout( + model_config, + mapping, + spec_config=spec_config, + layer_mask=draft_layer_mask, + ) + assert local_mamba_layers == 0 + assert local_attention_layers == rank + + draft_cost = V2MambaHybridCacheManager.get_cache_size_per_token( + model_config, + mapping, + max_batch_size=1, + kv_cache_config=KvCacheConfig(enable_block_reuse=False), + num_layers=1, + spec_config=spec_config, + layer_mask=draft_layer_mask, + ) + assert draft_cost == (64 * rank, 4096 * rank) + + +@pytest.mark.parametrize( + ("spec_config", "enable_attention_dp", "expected_intercept"), + [ + (None, False, 1728), + (MTPDecodingConfig(max_draft_len=4), True, 1792), + ( + MTPDecodingConfig( + max_draft_len=4, + draft_len_schedule={1: 4, 2: 2, 3: 1}, + ), + True, + 1984, + ), + ], +) +def test_v2_hybrid_estimator_accounts_for_ssm_slot_attention_bound( + monkeypatch, spec_config, enable_attention_dp, expected_intercept +): + monkeypatch.setattr( + KVCacheManager, + "get_cache_size_per_token", + lambda *args, **kwargs: 11, + ) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.mamba_cache_manager._get_local_mamba_cache_layout", + lambda *args, **kwargs: ( + SimpleNamespace(get_states_bytes_per_layer=lambda mapping: 64), + 1, + 1, + ), + ) + mapping = Mapping( + world_size=1, + tp_size=1, + pp_size=1, + enable_attention_dp=enable_attention_dp, + ) + + assert V2MambaHybridCacheManager.get_cache_size_per_token( + object(), + mapping, + max_batch_size=4, + kv_cache_config=KvCacheConfig(), + spec_config=spec_config, + ) == (12, expected_intercept) + + +def test_v2_hybrid_estimator_reserves_attention_for_each_lineage(monkeypatch): + monkeypatch.setattr( + KVCacheManager, + "get_cache_size_per_token", + lambda *args, **kwargs: 11, + ) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.mamba_cache_manager._get_local_mamba_cache_layout", + lambda *args, **kwargs: ( + SimpleNamespace(get_states_bytes_per_layer=lambda mapping: 64), + 1, + 1, + ), + ) + + # The one CUDA-graph padding slot is less than the four resident request + # lineages, but the conservative attention bound still reserves one page + # for every lineage. + assert V2MambaHybridCacheManager.get_cache_size_per_token( + object(), + Mapping(world_size=1, tp_size=1, pp_size=1), + max_batch_size=4, + kv_cache_config=KvCacheConfig(enable_block_reuse=False), + ) == (11, 1728) + + +def test_v2_hybrid_attention_bound_is_snapshot_alignment_agnostic(): + model_config = _hybrid_cache_sizing_model_config(["linear_attention", "full_attention"]) + mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + + def estimate(interval): + return V2MambaHybridCacheManager.get_cache_size_per_token( + model_config, + mapping, + max_batch_size=2, + kv_cache_config=KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=interval), + ), + tokens_per_block=32, + max_seq_len=128, + ) + + aligned = estimate(32) + unaligned = estimate(48) + + assert aligned[1] == unaligned[1] + + +def test_v2_hybrid_planned_capacity_is_bounded_by_resident_sequences(): + mgr = object.__new__(V2MambaHybridCacheManager) + mgr.max_batch_size = 4 + mgr.mapping = Mapping(world_size=2, rank=0, tp_size=1, pp_size=2) + mgr.max_seq_len = 128 + mgr._get_max_tokens_from_quota = lambda quota: 1 << 30 + + capacity = mgr._planned_token_capacity(KvCacheConfig(max_tokens=None), 1 << 40) + + assert capacity == 4 * 2 * 128 + + +def test_v2_hybrid_typical_batch_uses_num_ssm_slots(): + mgr = object.__new__(V2MambaHybridCacheManager) + mgr.kv_cache_type = CacheTypeCpp.SELF + mgr.head_dim_per_layer = [64, 64] + mgr.pp_layers = [0, 1] + mgr._mamba_layer_mask = [True, False] + mgr.ssm_bytes = 64 + mgr.conv_bytes = 32 + mgr.max_attention_window_vec = [128, 128] + mgr.max_batch_size = 2 + mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + mgr.max_seq_len = 128 + mgr.max_num_tokens = 128 + mgr.tokens_per_block = 32 + mgr.num_local_layers = 2 + mgr.local_num_mamba_layers = 1 + mgr._num_reserved_dummy_slots = 1 + mgr.dtype = DataType.HALF + mgr.enable_swa_scratch_reuse = False + mgr.enable_stats = False + mgr.num_extra_kv_tokens = 0 + mgr.get_layer_bytes_per_token = lambda **kwargs: 8 + mgr._minimum_live_gpu_quota = lambda: 0 + mgr._planned_token_capacity = lambda *_args: 128 + kv_cache_config = KvCacheConfig( + enable_partial_reuse=True, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=48), + ) + + config = mgr._build_cache_config( + kv_cache_config, + tokens_per_block=32, + vocab_size=None, + cache_tiers=[GpuCacheTierConfig(quota=1 << 20)], + ) + + assert len(config.typical_step.kv_caches) == 2 + assert [kv.num_ssm_slots for kv in config.typical_step.kv_caches] == [3, 2] + assert not hasattr(config.typical_step, "ssm_cache") + assert not hasattr(config.typical_step, "additional_attention_cache") + assert not hasattr(config.typical_step.kv_caches[0], "num_extra_attention_slots") + + +def test_v2_hybrid_rejects_quota_below_live_state_floor(): + mgr = object.__new__(V2MambaHybridCacheManager) + mgr.max_batch_size = 2 + mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + mgr.local_num_mamba_layers = 1 + mgr.ssm_bytes = 64 + mgr.conv_bytes = 32 + mgr._num_reserved_dummy_slots = 1 + mgr.tokens_per_block = 32 + mgr.num_local_layers = 2 + mgr.pp_layers = [0, 1] + mgr.max_attention_window_vec = [128, 128] + mgr.max_num_tokens = 128 + mgr.enable_swa_scratch_reuse = False + mgr.get_layer_bytes_per_token = lambda **kwargs: 8 + mgr._attention_cache_bytes_per_token = lambda: 16 + minimum_quota = mgr._minimum_live_gpu_quota() + + with pytest.raises(ValueError, match="too small for live recurrent states"): + mgr._build_cache_config( + KvCacheConfig(enable_partial_reuse=False), + tokens_per_block=32, + vocab_size=None, + cache_tiers=[GpuCacheTierConfig(quota=minimum_quota - 1)], + ) + + +def test_v2_hybrid_pure_mamba_rank_does_not_reserve_attention_page(): + mgr = object.__new__(V2MambaHybridCacheManager) + mgr.max_batch_size = 2 + mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + mgr.local_num_mamba_layers = 1 + mgr.ssm_bytes = 64 + mgr.conv_bytes = 32 + mgr._num_reserved_dummy_slots = 1 + mgr.tokens_per_block = 32 + mgr.num_local_layers = 1 + mgr.pp_layers = [0] + mgr.max_attention_window_vec = [128] + mgr.max_num_tokens = 128 + mgr.enable_swa_scratch_reuse = False + mgr.get_layer_bytes_per_token = lambda **kwargs: 0 + mgr._attention_cache_bytes_per_token = lambda: 0 + + assert mgr._minimum_live_gpu_quota() == 3 * (64 + 32) + + def test_cpp_hybrid_prepare_expect_snapshot_points(): mgr = object.__new__(CppMambaHybridCacheManager) mgr.kv_cache_config = KvCacheConfig( enable_block_reuse=True, - mamba_state_cache_interval=64, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=64), ) mgr.linear_attention_metadata = SimpleNamespace(states_snapshot_interval=64) requests = [ @@ -368,6 +1157,68 @@ def test_cpp_hybrid_prepare_expect_snapshot_points(): ] +def test_v2_block_reuse_commit_saves_ssm_snapshot_at_snapshot_point(): + mgr = object.__new__(V2MambaHybridCacheManager) + mgr.enable_block_reuse = True + mgr.is_draft = False + mgr._augment_tokens_for_block_reuse = lambda tokens, request, start, end: tokens[start:end] + mgr._mark_context_position_as_history = MagicMock() + + token_ids = list(range(150)) + request = SimpleNamespace( + prompt_len=150, + context_current_position=137, + context_remaining_length=13, + expect_snapshot_points=[137], + is_dummy_request=False, + is_dummy=False, + py_request_id=0, + get_tokens=lambda beam_idx: token_ids, + ) + kv_cache = SimpleNamespace( + num_committed_tokens=0, + commit=MagicMock(), + stop_committing=MagicMock(), + ) + + mgr.try_commit_blocks(request, kv_cache) + + kv_cache.commit.assert_called_once_with(token_ids[:137]) + kv_cache.stop_committing.assert_not_called() + mgr._mark_context_position_as_history.assert_called_once_with(request, kv_cache) + + # The remaining suffix advances request history but must not publish a + # second attention/SSM snapshot beyond the configured boundary. + kv_cache.num_committed_tokens = 137 + request.context_current_position = 150 + request.context_remaining_length = 0 + mgr.try_commit_blocks(request, kv_cache) + + kv_cache.commit.assert_called_once_with(token_ids[:137]) + kv_cache.stop_committing.assert_called_once_with() + + +def test_v2_hybrid_add_dummy_requests_forwards_encoder_output_lens(mocker): + mgr = object.__new__(V2MambaHybridCacheManager) + base_add_dummy_requests = mocker.patch.object( + KVCacheManagerV2, "add_dummy_requests", return_value=[] + ) + + mgr.add_dummy_requests([123], encoder_output_lens=[17]) + + assert base_add_dummy_requests.call_args.kwargs["encoder_output_lens"] == [17] + + +def test_v2_hybrid_prepare_expect_snapshot_points_clears_when_reuse_disabled(): + mgr = object.__new__(V2MambaHybridCacheManager) + mgr.kv_cache_config = KvCacheConfig(enable_block_reuse=False) + request = SimpleNamespace(prompt_len=150, expect_snapshot_points=[64]) + + mgr.prepare_expect_snapshot_points([request]) + + assert request.expect_snapshot_points == [] + + def test_cpp_hybrid_prepare_expect_snapshot_points_clears_when_reuse_disabled(): mgr = object.__new__(CppMambaHybridCacheManager) mgr.kv_cache_config = KvCacheConfig(enable_block_reuse=False) @@ -404,16 +1255,71 @@ def test_expect_snapshot_points_binding_round_trip(): assert request.expect_snapshot_points == [64, 128] +@skip_no_cuda +def test_v2_hybrid_pool_ratio_controls_allocated_memory(): + def allocated_memory(pool_ratio): + mgr = object.__new__(V2MambaHybridCacheManager) + mgr.kv_cache_type = CacheTypeCpp.SELF + mgr.head_dim_per_layer = [64, 64] + mgr.pp_layers = [0, 1] + mgr._mamba_layer_mask = [True, False] + mgr.ssm_bytes = 64 + mgr.conv_bytes = 32 + mgr.max_attention_window_vec = [128, 128] + mgr.max_batch_size = 2 + mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + mgr.max_seq_len = 128 + mgr.max_num_tokens = 128 + mgr.tokens_per_block = 32 + mgr.num_local_layers = 2 + mgr.local_num_mamba_layers = 1 + mgr._num_reserved_dummy_slots = 1 + mgr.dtype = DataType.HALF + mgr.enable_swa_scratch_reuse = False + mgr.enable_stats = False + mgr.num_extra_kv_tokens = 0 + mgr.get_layer_bytes_per_token = lambda **kwargs: 8 + + kv_cache_config = KvCacheConfig( + pool_ratio=pool_ratio, + enable_partial_reuse=False, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=64), + ) + mgr.kv_cache_config = kv_cache_config + config = mgr._build_cache_config( + kv_cache_config, + tokens_per_block=32, + vocab_size=1024, + cache_tiers=[GpuCacheTierConfig(quota=64 << 20)], + ) + runtime_manager = RuntimeKVCacheManager(config) + try: + statistics = runtime_manager._storage.get_statistics() + allocated_bytes = [ + int(stats.total) * sum(int(size) for size in stats.slot_size) + for stats in statistics + ] + return allocated_bytes, list(runtime_manager._current_gpu_ratio) + finally: + runtime_manager.shutdown() + + low_mamba_allocation, low_actual_ratio = allocated_memory([0.25, 0.75]) + high_mamba_allocation, high_actual_ratio = allocated_memory([0.75, 0.25]) + + assert low_actual_ratio == pytest.approx([0.25, 0.75]) + assert high_actual_ratio == pytest.approx([0.75, 0.25]) + assert high_mamba_allocation[0] > low_mamba_allocation[0] + assert high_mamba_allocation[1] < low_mamba_allocation[1] + + # --------------------------------------------------------------------------- -# CppMambaHybridCacheManager: recurrent-state snapshot pool sizing +# Cpp/V2 Mamba hybrid managers: recurrent-state allocation and reuse # -# Sized in KVCacheManager._calculate_max_num_blocks_for_linear_attention. -# Mirrors the MixedMambaCacheManager fix where each kind of padding sentinel -# (CUDA-graph dummy, plus one per draft length under spec decoding) must not -# evict live recurrent state. Wanli's fix made all sentinels share one slot -# in the Python manager (#13489); the C++ hybrid path instead reserves a -# dedicated slot per sentinel kind in the underlying pool — same invariant, -# different mechanism. These tests guard the pool sizing. +# The Cpp pool is sized in +# KVCacheManager._calculate_max_num_blocks_for_linear_attention. It reserves a +# dedicated slot for each padding sentinel kind so dummy requests cannot evict +# live recurrent state. The V2 tests cover unified-pool state views, slot +# bookkeeping, replay, and snapshot reuse. # --------------------------------------------------------------------------- @@ -421,12 +1327,13 @@ def _build_hybrid_with_mamba_layer( spec_config=None, max_batch_size=4, enable_block_reuse=False, - mamba_state_cache_interval=256, + periodic_snapshot_interval=256, is_estimating_kv_cache=False, dtype=DataType.HALF, mamba_layer_mask=None, attention_layer_mask=None, mamba_ssm_cache_dtype=torch.float16, + use_replay_state_update=False, ): """Construct a real CppMambaHybridCacheManager with one mamba layer + one full-attention layer so the parent KVCacheManager goes through the @@ -439,7 +1346,7 @@ def _build_hybrid_with_mamba_layer( kv_cache_config = KvCacheConfig( max_tokens=512, enable_block_reuse=enable_block_reuse, - mamba_state_cache_interval=mamba_state_cache_interval, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=periodic_snapshot_interval), ) return CppMambaHybridCacheManager( mamba_d_state=8, @@ -464,7 +1371,465 @@ def _build_hybrid_with_mamba_layer( layer_mask=attn_mask, is_estimating_kv_cache=is_estimating_kv_cache, dtype=dtype, + use_replay_state_update=use_replay_state_update, + ) + + +def _build_v2_hybrid_with_mamba_layer( + max_batch_size=4, + num_mamba_layers=1, + num_attention_layers=1, + num_kv_heads=4, + mapping=None, + spec_config=None, + use_replay_state_update=False, + enable_block_reuse=False, + enable_partial_reuse=True, + enable_attention_dp=False, + enable_swa_scratch_reuse=False, +): + """Construct a real V2MambaHybridCacheManager.""" + mamba_mask = [True] * num_mamba_layers + [False] * num_attention_layers + attn_mask = [False] * num_mamba_layers + [True] * num_attention_layers + if mapping is None: + mapping = Mapping( + world_size=1, + rank=0, + tp_size=1, + pp_size=1, + enable_attention_dp=enable_attention_dp, + ) + kv_cache_config = KvCacheConfig( + max_tokens=512, + enable_block_reuse=enable_block_reuse, + enable_partial_reuse=enable_partial_reuse, + enable_swa_scratch_reuse=enable_swa_scratch_reuse, + ) + return V2MambaHybridCacheManager( + mamba_d_state=8, + mamba_d_conv=4, + mamba_num_heads=4, + mamba_n_groups=1, + mamba_head_dim=8, + mamba_num_layers=num_mamba_layers, + mamba_layer_mask=mamba_mask, + mamba_cache_dtype=torch.float16, + mamba_ssm_cache_dtype=torch.float16, + kv_cache_config=kv_cache_config, + kv_cache_type=CacheTypeCpp.SELF, + num_layers=num_attention_layers, + num_kv_heads=num_kv_heads, + head_dim=64, + tokens_per_block=32, + max_seq_len=128, + max_batch_size=max_batch_size, + mapping=mapping, + spec_config=spec_config, + layer_mask=attn_mask, + vocab_size=1024, + use_replay_state_update=use_replay_state_update, + ) + + +def _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5): + """Spec config whose per-step token width is wider than draft depth. + + This mirrors parallel-draft style metadata closely enough for cache-manager + sizing without constructing the full speculative worker stack. + """ + return SimpleNamespace( + max_draft_len=max_draft_len, + max_total_draft_tokens=tokens_per_gen_step - 1, + tokens_per_gen_step=tokens_per_gen_step, + spec_dec_mode=SimpleNamespace(use_one_engine=lambda: False), + ) + + +def _assert_replay_layer_cache_uses_history_size(layer_cache, history_size): + assert layer_cache.old_x is not None + assert layer_cache.old_B is not None + assert layer_cache.old_dt is not None + assert layer_cache.old_dA_cumsum is not None + assert layer_cache.cache_buf_idx is not None + assert layer_cache.prev_num_accepted_tokens is not None + assert layer_cache.old_x.dim() == 5 + cache_size = layer_cache.temporal.shape[0] + assert layer_cache.old_x.shape[0] == cache_size + assert layer_cache.old_B.shape[0] == cache_size + assert layer_cache.old_dt.shape[0] == cache_size + assert layer_cache.old_dA_cumsum.shape[0] == cache_size + assert layer_cache.cache_buf_idx.shape[0] == cache_size + assert layer_cache.prev_num_accepted_tokens.shape[0] == cache_size + assert layer_cache.old_x.shape[1] == 2 + assert layer_cache.old_B.shape[1] == 2 + assert layer_cache.old_dt.shape[1] == 2 + assert layer_cache.old_dA_cumsum.shape[1] == 2 + assert layer_cache.old_x.shape[2] == history_size + assert layer_cache.old_B.shape[2] == history_size + assert layer_cache.old_dt.shape[-1] == history_size + assert layer_cache.old_dA_cumsum.shape[-1] == history_size + + +@skip_no_cuda +def test_v2_hybrid_allocates_mamba_state_and_dummy_indices(): + mgr = _build_v2_hybrid_with_mamba_layer(max_batch_size=4) + try: + assert mgr.local_num_mamba_layers == 1 + assert len(mgr.all_ssm_states) == 1 + assert len(mgr.all_conv_states) == 1 + assert mgr.all_ssm_states[0].shape[1:] == torch.Size([4, 8, 8]) + assert mgr.all_conv_states[0].shape[1:] == torch.Size([48, 3]) + assert mgr.get_max_resource_count() == 4 + assert mgr.blocks_in_primary_pool > 0 + assert isinstance(mgr.check_invalid_values_in_kv_cache(), bool) + + requests = mgr.add_dummy_requests([123], token_nums=[8], is_gen=False) + + assert len(requests) == 1 + indices = mgr.get_state_indices([123], [False]) + assert len(indices) == 1 + assert indices[0] >= 0 + assert mgr.cuda_state_indices[0].item() == indices[0] + assert mgr.get_ssm_states(0).data_ptr() == mgr.all_ssm_states[0].data_ptr() + assert mgr.get_conv_states(0).data_ptr() == mgr.all_conv_states[0].data_ptr() + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_supports_pure_mamba_pp_rank(): + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=2, + num_attention_layers=0, + ) + try: + assert mgr.local_num_mamba_layers == 1 + assert mgr.blocks_in_primary_pool == 0 + assert mgr._attention_cache_bytes_per_token() == 0 + assert len(mgr.all_ssm_states) == 1 + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_metadata_supports_attention_only_pp_rank(): + mapping = Mapping(world_size=2, rank=1, tp_size=1, pp_size=2) + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=1, + mapping=mapping, + ) + try: + assert mgr.local_num_mamba_layers == 0 + mgr.add_dummy_requests([123], token_nums=[8], is_gen=False) + + metadata = Mamba2Metadata(max_batch_size=1, chunk_size=8) + seq_lens = torch.tensor([8], dtype=torch.int32) + metadata.prepare( + SimpleNamespace( + seq_lens=seq_lens, + seq_lens_cuda=seq_lens.cuda(), + num_contexts=1, + num_ctx_tokens=8, + kv_cache_manager=mgr, + request_ids=[123], + kv_cache_params=SimpleNamespace( + num_cached_tokens_per_seq=torch.tensor([0], dtype=torch.int32) + ), + ) + ) + + assert metadata.state_indices[0].item() == 0 + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_invalid_check_scans_distinct_attention_pools(): + mgr = _build_v2_hybrid_with_mamba_layer( + num_attention_layers=2, + num_kv_heads=[1, 2, 4], + ) + try: + mgr.check_invalid_values_in_kv_cache(fill_with_zero=True) + first_attention_buffer = mgr.get_buffers(1) + second_attention_buffer = mgr.get_buffers(2) + assert first_attention_buffer.data_ptr() != second_attention_buffer.data_ptr() + + second_attention_buffer.flatten()[0] = torch.nan + + assert mgr.check_invalid_values_in_kv_cache() + finally: + mgr.shutdown() + + +@skip_no_cuda +@pytest.mark.parametrize("max_batch_size", [1, 4]) +def test_v2_hybrid_dummy_indices_keep_cuda_buffer_address(max_batch_size): + mgr = _build_v2_hybrid_with_mamba_layer(max_batch_size=max_batch_size, enable_attention_dp=True) + try: + stale_request = mgr.add_dummy_requests([123], token_nums=[8], is_gen=False)[0] + state_indices_ptr = mgr.cuda_state_indices.data_ptr() + + # Model a transfer-pending prior batch. The new ADP dummy is the only + # request participating in the next forward pass. + mgr.requests = [stale_request] * max_batch_size + new_requests = mgr.add_dummy_requests( + [ATTENTION_DP_DUMMY_REQUEST_ID], token_nums=[8], is_gen=False + ) + + assert len(new_requests) == 1 + assert len(mgr.requests) == max_batch_size + 1 + expected_capacity = max_batch_size + mgr._num_reserved_dummy_slots + assert mgr.cuda_state_indices.shape[0] == expected_capacity + assert mgr._host_state_indices.shape[0] == expected_capacity + assert mgr.cuda_state_indices.data_ptr() == state_indices_ptr + assert ( + mgr.cuda_state_indices[0].item() + == mgr.get_state_indices([ATTENTION_DP_DUMMY_REQUEST_ID], [False])[0] + ) + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_reserves_every_persistent_dummy_slot(): + spec_config = MTPDecodingConfig( + max_draft_len=4, + draft_len_schedule={1: 4, 2: 2, 3: 1}, + ) + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=4, + spec_config=spec_config, + enable_attention_dp=True, + ) + try: + runtime_draft_lengths = [4, 2, 1, 0] + cuda_graph_dummy_ids = [ + CUDA_GRAPH_DUMMY_REQUEST_ID - draft_len for draft_len in runtime_draft_lengths + ] + request_ids = [101, 102, 103, 104] + + assert mgr._num_cuda_graph_padding_dummy_slots == 4 + assert mgr._num_attention_dp_dummy_slots == 1 + assert mgr._num_reserved_dummy_slots == 5 + assert mgr.index_mapper.num_free_slots() == len(request_ids) + 5 + + assert ( + mgr.add_dummy_requests(request_ids, token_nums=[1] * len(request_ids), is_gen=False) + is not None + ) + for request_id, draft_len in zip(cuda_graph_dummy_ids, runtime_draft_lengths): + assert ( + mgr.add_dummy_requests( + [request_id], + is_gen=True, + max_num_draft_tokens=draft_len, + ) + is not None + ) + assert ( + mgr.add_dummy_requests([ATTENTION_DP_DUMMY_REQUEST_ID], token_nums=[1], is_gen=False) + is not None + ) + + all_request_ids = request_ids + cuda_graph_dummy_ids + [ATTENTION_DP_DUMMY_REQUEST_ID] + state_indices = mgr.get_state_indices(all_request_ids, [False] * len(all_request_ids)) + assert len(set(state_indices)) == len(all_request_ids) + assert mgr.index_mapper.num_free_slots() == 0 + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_free_resources_drops_stale_state_index_mapping(): + mgr = _build_v2_hybrid_with_mamba_layer() + try: + request = mgr.add_dummy_requests([123], token_nums=[8], is_gen=False)[0] + request_id = request.py_request_id + assert request_id in mgr._request_id_to_state_index + + mgr.requests.clear() + mgr.free_resources(request) + + assert request_id not in mgr._request_id_to_state_index + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_uses_upstream_min_snapshot_policy(): + mgr = _build_v2_hybrid_with_mamba_layer( + enable_block_reuse=True, + enable_partial_reuse=True, + ) + try: + assert mgr.block_reuse_policy is BlockReusePolicy.PER_REQUEST + assert mgr.kv_cache_config.enable_partial_reuse + assert mgr.kv_cache_manager_py_config.commit_min_snapshot + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_mamba_state_views_use_logical_slots(): + mgr = _build_v2_hybrid_with_mamba_layer(max_batch_size=4, num_mamba_layers=2) + try: + assert len(mgr.all_ssm_states) == 2 + assert len(mgr.all_conv_states) == 2 + + ssm_slots = mgr.all_ssm_states[0].shape[0] + conv_slots = mgr.all_conv_states[0].shape[0] + assert all(t.shape[0] == ssm_slots for t in mgr.all_ssm_states) + assert all(t.shape[0] == conv_slots for t in mgr.all_conv_states) + assert ssm_slots == conv_slots + + for local_layer_idx, ssm_state, conv_state in zip( + mgr.mamba_local_layer_ids, mgr.all_ssm_states, mgr.all_conv_states + ): + layer_id = LayerId(local_layer_idx) + ssm_scale = mgr.impl.get_page_index_scale(layer_id, MambaRole.SSM_STATE) + conv_scale = mgr.impl.get_page_index_scale(layer_id, MambaRole.CONV_STATE) + assert ssm_state.stride(0) == mgr.ssm_count * ssm_scale + assert conv_state.stride(0) == mgr.conv_count * conv_scale + assert ( + ssm_state.shape[0] + == ( + mgr.impl.get_page_index_upper_bound(layer_id, MambaRole.SSM_STATE) + + ssm_scale + - 1 + ) + // ssm_scale + ) + assert ( + conv_state.shape[0] + == ( + mgr.impl.get_page_index_upper_bound(layer_id, MambaRole.CONV_STATE) + + conv_scale + - 1 + ) + // conv_scale + ) + + mgr.add_dummy_requests([123, 456], token_nums=[8, 8], is_gen=False) + indices = mgr.get_state_indices([123, 456], [False, False]) + assert all(0 <= index < ssm_slots for index in indices) + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_swa_scratch_keeps_ssm_placeholder_rows(): + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=4, + enable_swa_scratch_reuse=True, ) + try: + request_id = 123 + mgr.add_dummy_requests([request_id], token_nums=[8], is_gen=False) + block_offsets = torch.zeros( + mgr.num_attention_op_pools, + 1, + 2, + mgr.max_blocks_per_seq, + dtype=torch.int32, + device="cuda", + ) + + mgr.copy_batch_block_offsets( + block_offsets, + [request_id], + beam_width=1, + num_contexts=1, + num_seqs=1, + ) + torch.cuda.synchronize() + + assert mgr.num_attention_op_pools == mgr.num_local_layers + assert mgr.kv_cache_pool_mapping.shape[0] == mgr.num_local_layers + assert mgr.kv_cache_pool_pointers[0, 0].item() != 0 + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_cpp_hybrid_replay_buffers_size_by_tokens_per_gen_step(): + spec_config = _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5) + mgr = _build_hybrid_with_mamba_layer( + spec_config=spec_config, + max_batch_size=4, + use_replay_state_update=True, + ) + try: + replay_metadata = mgr.get_replay_state_update_metadata() + assert mgr.use_replay_state_update is True + assert replay_metadata is not None + assert replay_metadata.replay_step_width == spec_config.tokens_per_gen_step + assert replay_metadata.replay_history_size == max( + MIN_REPLAY_HISTORY_SIZE, spec_config.tokens_per_gen_step + ) + layer_cache = mgr.mamba_layer_cache(0) + _assert_replay_layer_cache_uses_history_size( + layer_cache, replay_metadata.replay_history_size + ) + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_replay_buffers_size_by_tokens_per_gen_step(): + spec_config = _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5) + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=4, + spec_config=spec_config, + use_replay_state_update=True, + ) + try: + replay_metadata = mgr.get_replay_state_update_metadata() + assert mgr.use_replay_state_update is True + assert replay_metadata is not None + assert replay_metadata.replay_step_width == spec_config.tokens_per_gen_step + assert replay_metadata.replay_history_size == spec_config.tokens_per_gen_step + layer_cache = mgr.mamba_layer_cache(0) + _assert_replay_layer_cache_uses_history_size( + layer_cache, replay_metadata.replay_history_size + ) + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_replay_bookkeeping_matches_checkpoint_predicate(monkeypatch): + spec_config = _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5) + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=4, + spec_config=spec_config, + use_replay_state_update=True, + ) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.mamba_cache_manager._promote_mamba_state_triton", + lambda *args, **kwargs: None, + ) + try: + slot = torch.tensor([0], dtype=torch.int32, device="cuda") + attn_metadata = SimpleNamespace(num_seqs=1, num_contexts=0) + + mgr.update_mamba_states( + attn_metadata, + torch.tensor([3], dtype=torch.int32, device="cuda"), + state_indices=slot, + ) + assert mgr.prev_num_accepted_tokens[0].item() == 3 + assert mgr.cache_buf_idx[0].item() == 0 + + mgr.update_mamba_states( + attn_metadata, + torch.tensor([2], dtype=torch.int32, device="cuda"), + state_indices=slot, + ) + assert mgr.prev_num_accepted_tokens[0].item() == 2 + assert mgr.cache_buf_idx[0].item() == 1 + finally: + mgr.shutdown() @skip_no_cuda @@ -639,7 +2004,7 @@ def test_cpp_hybrid_recurrent_pool_floor_with_block_reuse(): """With block reuse enabled, the block-reuse branch must not drop the live-state + CUDA-graph-padding floor. - With max_batch_size=4, mamba_state_cache_interval=256, max_tokens=512: + With max_batch_size=4, periodic_snapshot_interval=256, max_tokens=512: naive: max_snapshots = 512 // 256 = 2 (drops live-state floor!) fixed: max_snapshots = max(2, 4 + 1) = 5 """ @@ -648,7 +2013,7 @@ def test_cpp_hybrid_recurrent_pool_floor_with_block_reuse(): spec_config=None, max_batch_size=max_batch_size, enable_block_reuse=True, - mamba_state_cache_interval=256, + periodic_snapshot_interval=256, ) recurrent_primary, _ = mgr.blocks_per_window[LinearCacheType.RECURRENT_STATES.value] assert recurrent_primary >= max_batch_size + 1, ( @@ -671,7 +2036,7 @@ def test_cpp_hybrid_dry_run_recurrent_pool_additive_with_block_reuse(): spec_config=None, max_batch_size=max_batch_size, enable_block_reuse=True, - mamba_state_cache_interval=256, + periodic_snapshot_interval=256, is_estimating_kv_cache=True, ) recurrent_primary, _ = mgr.blocks_per_window[LinearCacheType.RECURRENT_STATES.value] @@ -693,9 +2058,8 @@ def test_cpp_hybrid_dry_run_recurrent_pool_additive_with_block_reuse(): # - call the real parent KVCacheManager with the union layer_mask and # num_layers=num_layers (not mamba_num_layers + num_layers), # - skip allocating any mamba-only state, and -# - leave self.requests = [] so the guards on prepare_resources / -# update_mamba_states / _setup_state_indices can no-op without touching -# uninitialized state. +# - leave self.requests = [] so Mamba-only hooks and metadata preparation +# can no-op without touching uninitialized state. # # We exercise the same Python branch with world_size=1 (so the real C++ # KVCacheManager init doesn't need MPI) and a layer mask that contains zero @@ -705,7 +2069,7 @@ def test_cpp_hybrid_dry_run_recurrent_pool_additive_with_block_reuse(): def _build_zero_mamba_hybrid( enable_block_reuse=False, - mamba_state_cache_interval=256, + periodic_snapshot_interval=256, ): """Construct a real CppMambaHybridCacheManager whose this-rank slice has no mamba layers. world_size=1 / pp_size=1 keeps the real parent @@ -722,7 +2086,7 @@ def _build_zero_mamba_hybrid( kv_cache_config = KvCacheConfig( max_tokens=128, enable_block_reuse=enable_block_reuse, - mamba_state_cache_interval=mamba_state_cache_interval, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=periodic_snapshot_interval), ) mgr = CppMambaHybridCacheManager( @@ -756,11 +2120,11 @@ def _build_zero_mamba_hybrid( @skip_no_cuda def test_cpp_hybrid_zero_local_mamba_layers(): """End-to-end: real parent KVCacheManager + real early-exit. Verifies - early-exit invariants on the manager state AND that the three guarded - methods no-op without raising on uninitialized mamba-only state.""" + early-exit invariants on the manager state and that Mamba hooks and + metadata preparation do not touch uninitialized Mamba-only state.""" mgr = _build_zero_mamba_hybrid( enable_block_reuse=True, - mamba_state_cache_interval=64, + periodic_snapshot_interval=64, ) # Early-exit indicators. @@ -819,3 +2183,20 @@ def test_cpp_hybrid_zero_local_mamba_layers(): mgr.prepare_resources(empty_batch) # super() runs, then guard returns mgr.update_mamba_states(attn_metadata=None, num_accepted_tokens=None, state_indices=None) mgr._setup_state_indices() + + metadata = Mamba2Metadata(max_batch_size=1, chunk_size=8) + seq_lens = torch.tensor([8], dtype=torch.int32) + metadata.prepare( + SimpleNamespace( + seq_lens=seq_lens, + seq_lens_cuda=seq_lens.cuda(), + num_contexts=1, + num_ctx_tokens=8, + kv_cache_manager=mgr, + request_ids=[123], + kv_cache_params=SimpleNamespace( + num_cached_tokens_per_seq=torch.tensor([0], dtype=torch.int32) + ), + ) + ) + assert metadata.state_indices[0].item() == 0 diff --git a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py index 15cab5d360c0..ee5c23b5faef 100644 --- a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py +++ b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py @@ -23,7 +23,7 @@ _MLA_KV_CACHE_REUSE_SUPPORTED_SM_VERSIONS, ) from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType -from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig +from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig, ContextChunkingPolicy from tensorrt_llm.quantization import QuantAlgo @@ -162,7 +162,11 @@ def _make_llm_args(): enable_partial_reuse=False, tokens_per_block=32, max_attention_window=None, - mamba_state_cache_interval=1, + mamba_state_config=SimpleNamespace( + periodic_snapshot_interval=256, + additional_snapshot_offsets_from_start=[], + additional_snapshot_offsets_from_end=[], + ), ) scheduler_config = SimpleNamespace( context_chunking_policy=None, @@ -209,6 +213,8 @@ def _run_create_py_executor( enable_flash_mla=False, model_max_seq_len=128, enable_chunked_prefill=False, + is_hybrid_linear_model=False, + ctx_chunk_configs=None, ): """Execute create_py_executor with mocked dependencies and return MLA runtime flags. @@ -224,6 +230,8 @@ def _run_create_py_executor( enable_flash_mla: Whether to emulate the FlashMLA block-size override. model_max_seq_len: Effective sequence length reported by the model engine. enable_chunked_prefill: Whether to request MLA chunked prefill support. + is_hybrid_linear_model: Whether to emulate a hybrid linear model. + ctx_chunk_configs: Optional list that receives the executor chunk config. Returns: Tuple of (kv_cache_reuse_flag, runtime_cache_reuse_flag, @@ -262,7 +270,11 @@ def _run_create_py_executor( monkeypatch.setattr(py_executor_creator, "_adjust_torch_mem_fraction", lambda: None) monkeypatch.setattr(py_executor_creator, "log_memory_usage", lambda *args, **kwargs: None) monkeypatch.setattr(py_executor_creator, "is_mla", lambda _: True) - monkeypatch.setattr(py_executor_creator, "is_hybrid_linear", lambda _: False) + monkeypatch.setattr( + py_executor_creator, + "is_hybrid_linear", + lambda _: is_hybrid_linear_model, + ) monkeypatch.setattr(py_executor_creator, "get_sm_version", lambda: sm_version) monkeypatch.setattr(py_executor_creator, "KvCacheCreator", _DummyKvCacheCreator) @@ -287,6 +299,8 @@ def _create_model_engine(**kwargs): monkeypatch.setattr(py_executor_creator, "PyTorchModelEngine", _create_model_engine) def _create_py_executor_instance(**kwargs): + if ctx_chunk_configs is not None: + ctx_chunk_configs.append(kwargs["ctx_chunk_config"]) return _DummyPyExecutor( resources=kwargs["resources"], model_engine=kwargs["model_engine"], @@ -415,6 +429,19 @@ def test_explicit_transceiver_buffer_size_is_preserved(monkeypatch): assert config.max_tokens_in_buffer == 256 +def test_hybrid_force_chunk_uses_block_alignment_unit(monkeypatch): + ctx_chunk_configs = [] + _run_create_py_executor( + monkeypatch, + sm_version=90, + kv_cache_quant_algo=QuantAlgo.NO_QUANT, + is_hybrid_linear_model=True, + ctx_chunk_configs=ctx_chunk_configs, + ) + + assert ctx_chunk_configs == [(ContextChunkingPolicy.FORCE_CHUNK, 32)] + + @pytest.mark.parametrize("sm_version", _MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS) def test_mla_supported_configuration_preserves_chunked_prefill(monkeypatch, sm_version): """Verify every supported MLA SM preserves chunked prefill when requested.""" diff --git a/tests/unittest/_torch/modeling/test_modeling_multimodal.py b/tests/unittest/_torch/modeling/test_modeling_multimodal.py index 112b665dfb8c..9d31c6b810c8 100644 --- a/tests/unittest/_torch/modeling/test_modeling_multimodal.py +++ b/tests/unittest/_torch/modeling/test_modeling_multimodal.py @@ -640,7 +640,7 @@ def get_hybrid_kv_cache_manager( head_dim = text_config.hidden_size // text_config.num_attention_heads # CppMambaHybridCacheManager reads Pydantic-only fields - # (mamba_state_cache_interval, enable_block_reuse) so we have to + # (mamba_state_config, enable_block_reuse) so we have to # construct the llmapi.llm_args.KvCacheConfig here, not the C++ # bindings KvCacheConfig that the standard KVCacheManager path uses. kv_cache_config = PyKvCacheConfig(max_tokens=num_blocks * tokens_per_block) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index d69ed8c4de69..168ae0db3fae 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -2140,6 +2140,104 @@ def test_ssm_reuse_keeps_snapshots_from_multiple_commits(self) -> None: kv4.resume(stream) kv4.close() + def test_ssm_same_block_snapshots_support_monotonic_multi_turn_reuse(self) -> None: + cfg = self._make_ssm_config(tokens_per_block=32, enable_partial_reuse=True) + self.manager = KVCacheManager(cfg) + engine = FakeEngine(cfg) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + + prompt = [self.next_token() for _ in range(64)] + + for snapshot_length, expected_reuse in ((10, 0), (20, 10), (25, 20)): + kv_cache = self.manager.create_kv_cache(input_tokens=prompt[:snapshot_length]) + self.assertEqual(kv_cache.num_committed_tokens, expected_reuse) + kv_cache.resume(stream) + kv_cache.capacity = snapshot_length + engine.execute( + [ + Step( + kv_cache, + prompt[expected_reuse:snapshot_length], + prompt[:expected_reuse], + ) + ], + stream, + ) + kv_cache.history_length = snapshot_length + kv_cache.commit(prompt[expected_reuse:snapshot_length]) + kv_cache.close() + + exact = self.manager.create_kv_cache(input_tokens=prompt[:25]) + self.assertEqual(exact.num_committed_tokens, 25) + exact.resume(stream) + exact.capacity = 25 + exact.history_length = 25 + engine.execute([Step(exact, [], prompt[:25])], stream) + exact.close() + + def test_ssm_same_block_forks_only_reuse_safe_snapshots(self) -> None: + cfg = self._make_ssm_config(tokens_per_block=32, enable_partial_reuse=True) + self.manager = KVCacheManager(cfg) + engine = FakeEngine(cfg) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + prompt = [self.next_token() for _ in range(64)] + + source = self.manager.create_kv_cache() + source.resume(stream) + commit_start = 0 + for commit_end in (10, 20): + source.capacity = commit_end + chunk = prompt[commit_start:commit_end] + engine.execute([Step(source, chunk, prompt[:commit_start])], stream) + source.history_length = commit_end + source.commit(chunk) + commit_start = commit_end + source.close() + + # The retained 20-token state is in the future of a fork at token 15. + # Falling back to zero reuse is safe; reusing that state would corrupt + # the fork's SSM history. + early_fork = prompt[:15] + [self.next_token() for _ in range(25)] + early = self.manager.create_kv_cache(input_tokens=early_fork) + self.assertEqual(early.num_committed_tokens, 0) + early.resume(stream) + early.capacity = len(early_fork) + engine.execute([Step(early, early_fork, [])], stream) + early.history_length = len(early_fork) + engine.execute([Step(early, [], early_fork)], stream) + early.close() + + later_fork = prompt[:25] + [self.next_token() for _ in range(15)] + later = self.manager.create_kv_cache(input_tokens=later_fork) + self.assertEqual(later.num_committed_tokens, 20) + later.resume(stream) + later.capacity = len(later_fork) + engine.execute([Step(later, later_fork[20:], later_fork[:20])], stream) + later.history_length = len(later_fork) + engine.execute([Step(later, [], later_fork)], stream) + later.close() + + aligned = self.manager.create_kv_cache(input_tokens=prompt[:32]) + self.assertEqual(aligned.num_committed_tokens, 20) + aligned.resume(stream) + aligned.capacity = 32 + engine.execute([Step(aligned, prompt[20:32], prompt[:20])], stream) + aligned.history_length = 32 + aligned.commit(prompt[20:32]) + aligned.close() + + aligned_fork = prompt[:40] + [self.next_token() for _ in range(8)] + reused = self.manager.create_kv_cache(input_tokens=aligned_fork) + self.assertEqual(reused.num_committed_tokens, 32) + reused.resume(stream) + reused.capacity = len(aligned_fork) + engine.execute([Step(reused, aligned_fork[32:], aligned_fork[:32])], stream) + reused.history_length = len(aligned_fork) + engine.execute([Step(reused, [], aligned_fork)], stream) + reused.close() + def test_ssm_partial_snapshot_respects_partial_reuse_setting(self) -> None: """Partial SSM snapshots are created, but partial prompt reuse remains optional.""" tokens_per_block = 32 @@ -2377,6 +2475,9 @@ def tearDown(self) -> None: # Non-power-of-2 sizes so granularity rounding is non-trivial. PG0_SLOT_SIZE = 786432 # 768KB (windowed) PG1_SLOT_SIZE = 1310720 # 1280KB (non-windowed) + SSM_STATE_SLOT_SIZE = 23592960 + SSM_CONV_SLOT_SIZE = 829440 + ATTN_SLOT_SIZE = 245760 def _make_config( self, @@ -2431,6 +2532,38 @@ def _make_config( swa_scratch_reuse=(SwaScratchReuseConfig() if enable_swa_scratch_reuse else None), ) + def _make_hybrid_config(self, gpu_quota: int = 128 << 20) -> KVCacheManagerConfig: + return KVCacheManagerConfig( + tokens_per_block=self.TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(quota=gpu_quota)], + layers=[ + SsmLayerConfig( + layer_id=LayerId(0), + buffers=[ + BufferConfig( + role=DataRole("ssm_state"), + size=self.SSM_STATE_SLOT_SIZE, + ), + BufferConfig( + role=DataRole("conv_state"), + size=self.SSM_CONV_SLOT_SIZE, + ), + ], + ), + AttentionLayerConfig( + layer_id=LayerId(1), + buffers=[ + BufferConfig( + role=DataRole("key"), + size=self.ATTN_SLOT_SIZE, + ), + ], + ), + ], + enable_partial_reuse=False, + commit_min_snapshot=True, + ) + def test_default_init_ratio(self): """Without typical_step or constraints, uses hardcoded fallback.""" cfg = self._make_config() @@ -2468,6 +2601,101 @@ def test_typical_step_long_sequences(self): self.assertLess(ratio[0], 0.15) manager.shutdown() + def test_num_ssm_slots_controls_ssm_slot_count(self): + """SSM sizing uses explicit state-slot count, not token capacity.""" + manager = KVCacheManager(self._make_hybrid_config()) + ssm_lc = manager._life_cycles.ssm_life_cycle_id + assert ssm_lc is not None + ssm_pg = manager._storage.get_pool_group_index(ssm_lc) + attn_pg = 1 - ssm_pg + + batch = BatchDesc( + kv_caches=[ + KVCacheDesc( + capacity=1024, + history_length=1023, + num_ssm_slots=3, + ) + ] + * 2 + ) + slots = manager._storage._compute_slots_for_batch(batch, self.TOKENS_PER_BLOCK, None) + self.assertEqual(slots[ssm_pg], 6) + self.assertEqual( + slots[attn_pg], + 2 * div_up(1024, self.TOKENS_PER_BLOCK) + 2, + ) + + default_batch = BatchDesc(kv_caches=[KVCacheDesc(capacity=4096, history_length=4095)] * 2) + default_slots = manager._storage._compute_slots_for_batch( + default_batch, self.TOKENS_PER_BLOCK, None + ) + self.assertEqual(default_slots[ssm_pg], 2) + manager.shutdown() + + def test_num_ssm_slots_must_be_positive(self): + with self.assertRaises(AssertionError): + KVCacheDesc(capacity=1, history_length=0, num_ssm_slots=0) + + def test_ssm_slots_bound_additional_attention_capacity(self): + manager = KVCacheManager(self._make_hybrid_config()) + ssm_lc = manager._life_cycles.ssm_life_cycle_id + assert ssm_lc is not None + ssm_pg = manager._storage.get_pool_group_index(ssm_lc) + attn_pg = 1 - ssm_pg + batch = BatchDesc([KVCacheDesc(capacity=64, history_length=63, num_ssm_slots=3)]) + + slots = manager._storage._compute_slots_for_batch(batch, self.TOKENS_PER_BLOCK, None) + + # Two SSM slots beyond the live state imply at most one retained + # partial attention page for this request lineage. + self.assertEqual(slots[attn_pg], 3) + manager.shutdown() + + def test_ssm_slot_attention_bound_applies_to_each_lifecycle(self): + config = self._make_config(enable_swa_scratch_reuse=True) + manager = KVCacheManager(config) + base_batch = BatchDesc([KVCacheDesc(capacity=512, history_length=32)]) + snapshot_batch = BatchDesc([KVCacheDesc(capacity=512, history_length=32, num_ssm_slots=3)]) + + base_slots = manager._storage._compute_slots_for_batch( + base_batch, self.TOKENS_PER_BLOCK, config.swa_scratch_reuse + ) + snapshot_slots = manager._storage._compute_slots_for_batch( + snapshot_batch, self.TOKENS_PER_BLOCK, config.swa_scratch_reuse + ) + expected_increase = [0] * manager._storage.num_pool_groups + for life_cycle_id, _ in manager._life_cycles.attention_life_cycles(): + pool_group = manager._storage.get_pool_group_index(life_cycle_id) + expected_increase[pool_group] += 1 + + self.assertEqual( + [actual - base for actual, base in zip(snapshot_slots, base_slots)], + expected_increase, + ) + manager.shutdown() + + def test_ssm_slot_attention_bound_reserves_each_lineage(self): + manager = KVCacheManager(self._make_hybrid_config()) + ssm_lc = manager._life_cycles.ssm_life_cycle_id + assert ssm_lc is not None + ssm_pg = manager._storage.get_pool_group_index(ssm_lc) + attn_pg = 1 - ssm_pg + batch = BatchDesc( + [ + KVCacheDesc(capacity=64, history_length=63, num_ssm_slots=3), + *[KVCacheDesc(capacity=64, history_length=63)] * 3, + ] + ) + + slots = manager._storage._compute_slots_for_batch(batch, self.TOKENS_PER_BLOCK, None) + + self.assertEqual(slots[ssm_pg], 6) + # Any non-live SSM capacity enables the conservative upper bound of one + # partial attention page for every request lineage. + self.assertEqual(slots[attn_pg], 12) + manager.shutdown() + def test_constraints_floor_typical_step(self): """Constraints clamp the typical_step ratio from below.""" typical = BatchDesc(kv_caches=[KVCacheDesc(capacity=4096, history_length=4000)] * 32) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py index a43818a18090..9c130f12977e 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py @@ -60,6 +60,7 @@ class _StatsRequest: state: LlmRequestState = LlmRequestState.GENERATION_IN_PROGRESS context_current_position: int = 0 context_chunk_size: int = 0 + expect_snapshot_points: list[int] = field(default_factory=list) prepopulated_prompt: tuple[int, int] | None = None kv_cache_perf_metric_calls: list[dict[str, int]] = field(default_factory=list) multimodal_hashes: None = None diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 4f17a87123c3..807d54baeff9 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -18,6 +18,7 @@ from utils.util import force_ampere import tensorrt_llm.bindings.executor as tle +import tensorrt_llm.llmapi as public_llmapi import tensorrt_llm.llmapi.llm_args as llm_args_mod from tensorrt_llm import LLM as TorchLLM from tensorrt_llm._torch.auto_deploy.llm_args import \ @@ -42,7 +43,8 @@ ExecutorMemoryType, ExtendedRuntimePerfKnobConfig, KvCacheConfig, - LookaheadDecodingConfig, MoeConfig, + LookaheadDecodingConfig, + MambaStateConfig, MoeConfig, MTPDecodingConfig, MultimodalConfig, MultimodalEncoderCudaGraphConfig, PeftCacheConfig, PybindMirror, @@ -133,6 +135,22 @@ def test_llm_args_with_kvcache_config(self): 1024, 1024, 1024 ] + def test_from_yaml_migrates_legacy_mamba_interval(self, tmp_path): + yaml_path = tmp_path / "legacy_mamba_interval.yaml" + yaml_path.write_text( + yaml.safe_dump({ + "model": str(llama_model_path), + "kv_cache_config": { + "mamba_state_cache_interval": 64, + }, + }), + encoding="utf-8", + ) + + llm_args = TorchLlmArgs.from_yaml(yaml_path) + + assert llm_args.kv_cache_config.mamba_state_config.periodic_snapshot_interval == 64 + def test_llm_args_with_pydantic_options(self): yaml_content = """ max_batch_size: 16 @@ -564,6 +582,11 @@ def test_KvCacheConfig_declaration(): enable_swa_scratch_reuse=True, enable_partial_reuse=True, copy_on_partial_reuse=True, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=0, + additional_snapshot_offsets_from_start=[128], + additional_snapshot_offsets_from_end=[0, 32], + ), pool_ratio=[0.25, 0.75], avg_seq_len=2048, block_reuse_policy="per_request", @@ -586,6 +609,13 @@ def test_KvCacheConfig_declaration(): assert config.pool_ratio == [0.25, 0.75] assert config.avg_seq_len == 2048 assert config.block_reuse_policy == "per_request" + assert config.mamba_state_config.periodic_snapshot_interval == 0 + assert config.mamba_state_config.additional_snapshot_offsets_from_start == [ + 128 + ] + assert config.mamba_state_config.additional_snapshot_offsets_from_end == [ + 0, 32 + ] assert not hasattr(pybind_config, "pool_ratio") assert not hasattr(pybind_config, "avg_seq_len") assert not hasattr(pybind_config, "block_reuse_policy") @@ -607,6 +637,92 @@ def test_KvCacheConfig_declaration(): KvCacheConfig(block_reuse_policy="invalid") +def test_MambaStateConfig_defaults_use_independent_lists(): + first = MambaStateConfig() + second = MambaStateConfig() + + assert first.periodic_snapshot_interval == 256 + first.additional_snapshot_offsets_from_start.append(128) + first.additional_snapshot_offsets_from_end.append(0) + + assert second.additional_snapshot_offsets_from_start == [] + assert second.additional_snapshot_offsets_from_end == [] + assert public_llmapi.MambaStateConfig is MambaStateConfig + + +def test_MambaStateConfig_rejects_unknown_fields(): + with pytest.raises(ValidationError, match="extra_forbidden"): + MambaStateConfig(unknown_snapshot_policy=1) + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("periodic_snapshot_interval", -1), + ("additional_snapshot_offsets_from_start", [0]), + ("additional_snapshot_offsets_from_start", [-1]), + ("additional_snapshot_offsets_from_start", [True]), + ("additional_snapshot_offsets_from_start", [1.5]), + ("additional_snapshot_offsets_from_start", ["128"]), + ("additional_snapshot_offsets_from_end", [-1]), + ("additional_snapshot_offsets_from_end", [True]), + ("additional_snapshot_offsets_from_end", [1.5]), + ("additional_snapshot_offsets_from_end", ["32"]), + ], +) +def test_MambaStateConfig_rejects_invalid_snapshot_offsets(field, value): + with pytest.raises(ValidationError, match=field): + MambaStateConfig(**{field: value}) + + +def test_KvCacheConfig_rejects_removed_mamba_interval_field(): + with pytest.raises(ValidationError, match="extra_forbidden"): + KvCacheConfig(mamba_state_cache_interval=64) + + with pytest.raises(ValidationError, match="extra_forbidden"): + TorchLlmArgs( + model=llama_model_path, + kv_cache_config={"mamba_state_cache_interval": 64}, + ) + + +def test_config_file_merge_migrates_legacy_mamba_interval_without_mutating_input( +): + yaml_dict = { + "kv_cache_config": { + "mamba_state_cache_interval": 64, + "mamba_state_config": { + "additional_snapshot_offsets_from_end": [0], + }, + }, + } + + merged = update_llm_args_with_extra_dict({"model": "dummy"}, yaml_dict) + + assert merged[ + "kv_cache_config"].mamba_state_config.periodic_snapshot_interval == 64 + assert merged[ + "kv_cache_config"].mamba_state_config.additional_snapshot_offsets_from_end == [ + 0 + ] + assert yaml_dict["kv_cache_config"]["mamba_state_cache_interval"] == 64 + + +def test_config_file_merge_rejects_legacy_and_new_mamba_intervals(): + with pytest.raises(ValueError, match="cannot set both"): + update_llm_args_with_extra_dict( + {"model": "dummy"}, + { + "kv_cache_config": { + "mamba_state_cache_interval": 64, + "mamba_state_config": { + "periodic_snapshot_interval": 64, + }, + }, + }, + ) + + def test_KvCacheConfig_disk_cache_validation(tmp_path): config = KvCacheConfig(disk_cache_size=2048, disk_cache_path=str(tmp_path)) From f3d8f6a86aecfea645e7b9aa4f20493b6abfc1bb Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:12:00 +0800 Subject: [PATCH 02/22] [TRTLLM-11875][refactor] Simplify V2 Mamba cache management (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- PR3_MAMBA_V2_REVIEW.md | 239 ---- .../_torch/pyexecutor/kv_cache_manager_v2.py | 143 +- .../_torch/pyexecutor/mamba_cache_manager.py | 1263 +++++------------ tensorrt_llm/llmapi/llm_args.py | 12 + .../executor/test_mamba_cache_manager.py | 136 +- tests/unittest/llmapi/test_llm_args.py | 34 + 6 files changed, 590 insertions(+), 1237 deletions(-) delete mode 100644 PR3_MAMBA_V2_REVIEW.md diff --git a/PR3_MAMBA_V2_REVIEW.md b/PR3_MAMBA_V2_REVIEW.md deleted file mode 100644 index e24ddad92013..000000000000 --- a/PR3_MAMBA_V2_REVIEW.md +++ /dev/null @@ -1,239 +0,0 @@ - - -# PR3 V2 Mamba snapshot reuse review - -This document records the design decisions, code changes, experiments, and -verification for PR3. It is intended to make the final PR diff reviewable -against upstream `main`, which already contains PR1, rather than against an -older stacked-PR base. - -## Baseline - -- PR1 was squash-merged into upstream `main` as `bb90835f8a`. -- PR3 was rebuilt on the latest upstream `main` from its four PR3-specific - commits. -- The original PR1 commit chain and the two stacked-branch synchronization - merges were removed from PR3 history. - -The final PR3 diff must be reviewed relative to upstream `main`. - -## Review requirements - -1. Select the V1 Mamba manager by default. Select V2 only when - `use_kv_cache_manager_v2: true` is explicitly configured. -2. Move Mamba snapshot policy under `mamba_state_config`: - - `periodic_snapshot_interval` controls regular snapshot boundaries. - - `additional_snapshot_offsets_from_start` expresses fixed boundaries - measured from the prompt start. - - `additional_snapshot_offsets_from_end` expresses fixed boundaries - measured backward from the prompt end; zero selects the prompt end. -3. Reject unsupported combinations before manager construction: - - V1 supports periodic snapshots only. - - V2 does not support disaggregated serving. -4. Keep generic KV cache manager changes minimal and express Mamba behavior - through a specialized extension wherever possible. -5. Determine whether `commit_min_snapshot` removes the need for - `allow_prefix_sibling`, using multi-turn and forked-prefix reuse tests. -6. Explain and justify the extra attention-slot bound. - -## Configuration and manager-selection decisions - -`forward` and `backward` were rejected as field names because they can be -misread as execution or traversal directions. The selected names state both -the unit and reference anchor explicitly. - -| Serving mode | `use_kv_cache_manager_v2` | Additional offsets | Result | -| --- | --- | --- | --- | -| Aggregated | false or auto | Empty | V1 C++ manager | -| Aggregated | true | Any valid policy | V2 Mamba manager | -| Aggregated | false or auto | Non-empty | Configuration error | -| Disaggregated | false or auto | Empty | V1 C++ or mixed manager, according to the transceiver | -| Disaggregated | true | Any | Configuration error | -| Disaggregated | false or auto | Non-empty | Configuration error | - -The `TLLM_MAMBA_MANAGER_PREFERENCE=V2` override conflicts with the requirement -that V2 be selected only by the explicit public setting. The V2 override has -therefore been removed. An explicit V2 setting that conflicts with the CPP, -MIXED, or `TRTLLM_USE_PY_MAMBA` compatibility overrides is rejected rather -than silently routed to another implementation. - -The snapshot resolver uses these exact semantics: - -- `periodic_snapshot_interval=0` disables periodic snapshots; -- start offsets are strictly positive absolute token boundaries; -- end offsets are non-negative and resolve to `prompt_len - offset`, so zero - selects the final prompt boundary; -- resolved positions outside `(0, prompt_len]` are ignored; -- periodic and fixed boundaries are deduplicated and sorted. - -## Public configuration API - -This PR deliberately makes the snapshot-policy rename a breaking Python API -change. The former top-level `mamba_state_cache_interval` constructor argument -and Python property are removed without a model validator or deprecated alias, -and strict model construction rejects the removed field. YAML/JSON file loaders -retain a narrow compatibility migration to -`mamba_state_config.periodic_snapshot_interval`; setting both paths is an -error. All in-tree callers, examples, telemetry metadata, and canonical -documentation use the nested field. - -## KV cache extension audit - -The completed read-only audit found that the generic manager should retain only -two neutral PR3 changes: - -- a neutral constructor input for reserved index slots, used only to size the - stable `IndexMapper` storage; -- reporting every physical pool group in iteration statistics, which is a - generic correctness fix and needs a non-Mamba test. - -The remaining changes are Mamba-specific and have moved into -`V2MambaHybridCacheManager`, following `DeepseekV4CacheManager`: - -- CUDA-graph and attention-DP dummy-slot calculation; -- SSM/conv roles and page-table layout; -- SWA scratch layout for recurrent pools; -- snapshot commit, history, final-free, and context-update lifecycle; -- recurrent-state invalid-value checking. - -`KVCacheDesc.capacity` remains token capacity. The only Mamba-specific sizing -field is `num_ssm_slots`, which counts the live, snapshot, and dummy recurrent -state slots associated with that logical request. The storage planner sums it -for SSM lifecycles. If the batch has any non-live SSM capacity, the planner also -reserves one retained partial attention page per request lineage. There is no -separate `num_extra_attention_slots` field or layer-specific `BatchDesc` -descriptor. - -The subclass computes the exact CUDA-graph and attention-DP dummy count before -calling the base constructor, then passes that count through the neutral -`num_reserved_index_slots` extension. Its CUDA and pinned-host state-index -tensors are allocated once for `max_batch_size + reserved slots`; later dummy -insertion therefore cannot replace an aliased tensor or invalidate a captured -CUDA graph. The subclass also owns the mixed-pool page table, event-window -mapping for `SsmLayerConfig`, snapshot commit/history lifecycle, and state -diagnostics. The generic manager no longer knows any Mamba roles or policies. - -Metadata preparation still runs on attention-only pipeline-parallel ranks even -though no local Mamba kernel consumes state indices there. Both V1 and V2 now -return harmless zero placeholders on those ranks instead of consulting Mamba -state maps or buffers that are intentionally absent. The V2 invalid-value -diagnostic also scans every attention layer: a runtime layer group represents -a lifecycle, not necessarily one physical pool, so layers with different -buffer sizes cannot safely be deduplicated by group ID. - -The audit classified each generic KVCacheManagerV2 hunk as one of: - -- generic prerequisite already supplied by `commit_min_snapshot`; -- generic extension point required by Mamba; -- Mamba-only policy that belongs in `mamba_cache_manager.py`; -- obsolete code removable from PR3. - -## `allow_prefix_sibling` experiment - -The PR3-specific constructor flag, sibling-retention rule, alternate SSM -matcher, and both call sites were removed. `commit_min_snapshot` remains the -only runtime prerequisite. - -The first experiment exposed a generic replacement-order bug: when a longer -partial block replaced the only child of a root, `detach_next()` could prune -the now-empty root before the replacement was attached. The replacement is now -registered before covered children are detached. - -The resulting semantics are deliberately simpler: - -- monotonically extending snapshots at 10, 20, and 25 tokens reuse 10, 20, - and 25 tokens respectively; -- a fork sharing only 15 tokens cannot consume the retained 20-token state and - safely re-prefills from zero; -- a fork sharing 25 tokens reuses the 20-token state; -- a block-aligned 32-token snapshot remains reusable by a later fork. - -Only the latest partial snapshot within one token block and request lineage is -retained. An earlier same-block fork may therefore miss a reusable state, but -it cannot read a state computed from tokens beyond the fork boundary. This is -an accepted efficiency tradeoff for removing the custom sibling structure. - -## Extra attention slots - -These are physical attention-cache slots used -when a non-block-aligned Mamba snapshot preserves a copied partial attention -page in the radix tree. They are not request slots or CUDA-graph/attention-DP -dummy slots. Their capacity is additional to the normal attention-page budget -because the active request still owns a writable partial page while the radix -tree retains the snapshot copy. - -Each logical request descriptor budgets at least one live SSM state. The -planner uses `total_num_ssm_slots > num_requests` only to detect that the batch -has non-live SSM capacity. It then reserves one retained partial attention page -for every request lineage and adds that bound to every attention lifecycle; -nothing is supplied by an external caller. Covered partial siblings are removed -as a lineage advances, so more than one retained page per lineage is -unnecessary. - -This intentionally trades some capacity precision for a smaller interface. -The bound does not require `total_num_ssm_slots >= 2 * num_requests`. When the -number of non-live slots is smaller than the number of lineages, or when those -slots are aligned snapshots or reserved dummy states, it simply over-reserves -attention capacity. It never depends on snapshot alignment or slot -distribution. The normal request token capacity already accounts for a resumed -request's writable page; this bound covers the retained radix-tree snapshot -page. The V1 estimator retains its existing policy because this single-field -bound is specific to V2. - -## PP/spec layout and affine sizing - -One-model MTP with pipeline parallelism is expected to provide an explicit -`pp_partition` covering the base-model layers. The existing PP helper uses that -partition as the base/spec boundary and assigns appended speculative layers to -the last rank. Supporting an automatically derived boundary when -`pp_partition` is absent is intentionally left to a separate change. - -The hybrid affine estimator uses the same normalized masks and explicit PP -layout to count local attention and Mamba layers. This avoids two failures: - -- custom `pp_partition` being applied to an attention-only layer count and - raising during startup; -- a separate, attention-only MTP draft cache being charged for phantom target - Mamba state. - -Target-only and draft-only masks are shared between estimation and runtime -construction, so the budget split uses the physical layout each manager will -actually allocate. - -## Verification log - -- `python3 -m py_compile` passed after the configuration, routing, subclass, - and test migrations. -- Focused configuration/routing/snapshot-point tests: 34 passed. -- Full `test_mamba_cache_manager.py`: 85 passed. -- Three focused final-audit regressions passed: V1 and V2 metadata preparation - on ranks with no local Mamba layers, and invalid-value detection in a second - differently sized attention pool that shares a lifecycle. -- Six focused PP/MTP sizing cases passed across default/custom PP partitions, - Cpp/V2 estimators, target-only masks, and attention-only draft masks. -- The combined Mamba, cache-budget split, dual-pool, and executor-creator suite - passed: 151 tests. -- Generic KVCMv2 and executor-creator tests: 24 passed. -- Full `test_llm_args.py`: 255 passed, 3 skipped. A preceding attempt failed - because the environment's MPI session-reuse worker did not report an identity - within 60 seconds on either spawn attempt (`0/1` workers); it did not reach - Mamba configuration or manager code. -- The C++ estimator regression for disabled periodic snapshots on a - recurrent-only rank passed, as did the V2 partial-attention sizing test. -- Runtime radix-tree focused tests: 2 passed. -- Runtime `TestSSMSupport`: 13 passed. -- Runtime `TestNoBatching`: 24 passed, 12 skipped. -- `scripts/generate_llm_args_golden_manifest.py --check` passed. The manifest - exposes only `kv_cache_config.mamba_state_config.periodic_snapshot_interval`. -- The telemetry generator cannot load legacy `TrtLlmArgs` in this local - environment and would incorrectly erase that entire table. The generated - noise was discarded; the two existing Mamba telemetry rows were updated to - the manifest's canonical path without changing unrelated rows. -- Isort/YAPF formatting and the completed runtime experiment's lint/diff - checks passed. The final read-only audit found no remaining correctness - blockers after the PP/spec sizing fixes. -- Final manifest, syntax, diff, and full pre-commit checks passed before - publication. diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 8ef4967d1746..c3960b4774ea 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -1114,26 +1114,44 @@ def append_to_kv_heads_per_layer( self._log_kv_cache_pool_lifecycle_mapping() + def _get_pool_page_index_role(self, pool_id: int) -> DataRole: + return Role.KEY + + def _get_pool_paired_role(self, pool_id: int) -> Optional[DataRole]: + if self.kv_cache_type == CacheTypeCpp.SELFKONLY: + return None + return Role.VALUE + + def _get_block_scale_role_for_pool(self, pool_id: int) -> Optional[DataRole]: + if self.dtype != DataType.NVFP4 or self._get_pool_page_index_role(pool_id) != Role.KEY: + return None + return Role.KEY_BLOCK_SCALE + def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: kv_cache_pool_pointers_list = [] kv_cache_pool_mapping_list = [] block_scale_pool_pointers_list = [] if self.enable_swa_scratch_reuse: for layer_id in typed_range(LayerId(self.num_local_layers)): + pool_id = self.impl.get_layer_group_id(layer_id) + page_index_role = self._get_pool_page_index_role(pool_id) kv_cache_pool_pointers_list.append( [ self.impl.get_mem_pool_base_address( - layer_id, Role.KEY, PageIndexMode.PER_LAYER + layer_id, page_index_role, PageIndexMode.PER_LAYER ), 0, ] ) if self.dtype == DataType.NVFP4: + block_scale_role = self._get_block_scale_role_for_pool(pool_id) block_scale_pool_pointers_list.append( [ self.impl.get_mem_pool_base_address( - layer_id, Role.KEY_BLOCK_SCALE, PageIndexMode.PER_LAYER - ), + layer_id, block_scale_role, PageIndexMode.PER_LAYER + ) + if block_scale_role is not None + else 0, 0, ] ) @@ -1141,63 +1159,63 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: else: for pool_id in range(self.num_pools): layer_id = self.impl.layer_grouping[pool_id][0] + page_index_role = self._get_pool_page_index_role(pool_id) kv_cache_pool_pointers_list.append( [ self.impl.get_mem_pool_base_address( - layer_id, Role.KEY, PageIndexMode.SHARED + layer_id, page_index_role, PageIndexMode.SHARED ), 0, ] ) if self.dtype == DataType.NVFP4: + block_scale_role = self._get_block_scale_role_for_pool(pool_id) block_scale_pool_pointers_list.append( [ self.impl.get_mem_pool_base_address( - layer_id, Role.KEY_BLOCK_SCALE, PageIndexMode.SHARED - ), + layer_id, block_scale_role, PageIndexMode.SHARED + ) + if block_scale_role is not None + else 0, 0, ] ) for layer_id in typed_range(LayerId(self.num_local_layers)): layer_group_id = self.impl.get_layer_group_id(layer_id) - if self.dtype != DataType.NVFP4: - key_base_addr = kv_cache_pool_pointers_list[layer_group_id][0] - addr_offset = ( - self.impl.get_mem_pool_base_address( - layer_id, Role.KEY, PageIndexMode.SHARED - ) - - key_base_addr + page_index_role = self._get_pool_page_index_role(layer_group_id) + index_base_addr = kv_cache_pool_pointers_list[layer_group_id][0] + addr_offset = ( + self.impl.get_mem_pool_base_address( + layer_id, page_index_role, PageIndexMode.SHARED ) + - index_base_addr + ) + offset_divisor = self.impl.get_page_stride(layer_id, page_index_role) + if self._get_pool_paired_role(layer_group_id) is not None: + offset_divisor *= self.kv_factor + offset = exact_div(addr_offset, offset_divisor) + + if self.dtype != DataType.NVFP4: + block_scale_offset = None else: - key_base_addr = kv_cache_pool_pointers_list[layer_group_id][0] - block_scale_base_addr = block_scale_pool_pointers_list[layer_group_id][0] - addr_offset = ( - self.impl.get_mem_pool_base_address( - layer_id, Role.KEY, PageIndexMode.SHARED - ) - - key_base_addr - ) - block_scale_addr_offset = ( - self.impl.get_mem_pool_base_address( - layer_id, Role.KEY_BLOCK_SCALE, PageIndexMode.SHARED + block_scale_role = self._get_block_scale_role_for_pool(layer_group_id) + if block_scale_role is None: + block_scale_offset = None + else: + block_scale_base_addr = block_scale_pool_pointers_list[layer_group_id][0] + block_scale_addr_offset = ( + self.impl.get_mem_pool_base_address( + layer_id, block_scale_role, PageIndexMode.SHARED + ) + - block_scale_base_addr ) - - block_scale_base_addr - ) - block_scale_offset = exact_div( - block_scale_addr_offset, - self.get_layer_bytes_per_token(layer_id, Role.KEY_BLOCK_SCALE) - * self.kv_factor - * self.tokens_per_block, - ) - offset = exact_div( - addr_offset, - self.get_layer_bytes_per_token(layer_id, Role.KEY) - * self.kv_factor - * self.tokens_per_block, - ) + block_scale_divisor = self.impl.get_page_stride(layer_id, block_scale_role) + if self._get_pool_paired_role(layer_group_id) is not None: + block_scale_divisor *= self.kv_factor + block_scale_offset = exact_div(block_scale_addr_offset, block_scale_divisor) - if self.dtype == DataType.NVFP4: + if block_scale_offset is not None: assert block_scale_offset == offset, ( "Block scale offset and offset should be the same" ) @@ -1232,12 +1250,16 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: ) for pool_id in range(self.num_pools): layer_id = self.impl.layer_grouping[pool_id][0] - self.index_scales[pool_id] = self.impl.get_page_index_scale(layer_id, Role.KEY) - if self.kv_cache_type != CacheTypeCpp.SELFKONLY: + page_index_role = self._get_pool_page_index_role(pool_id) + self.index_scales[pool_id] = self.impl.get_page_index_scale(layer_id, page_index_role) + paired_role = self._get_pool_paired_role(pool_id) + if paired_role is not None: self.kv_offset[pool_id] = exact_div( - self.impl.get_mem_pool_base_address(layer_id, Role.VALUE, PageIndexMode.SHARED) - - self.impl.get_mem_pool_base_address(layer_id, Role.KEY, PageIndexMode.SHARED), - self.impl.get_page_stride(layer_id, Role.KEY), + self.impl.get_mem_pool_base_address(layer_id, paired_role, PageIndexMode.SHARED) + - self.impl.get_mem_pool_base_address( + layer_id, page_index_role, PageIndexMode.SHARED + ), + self.impl.get_page_stride(layer_id, page_index_role), ) else: self.kv_offset[pool_id] = 0 @@ -1355,7 +1377,8 @@ def _get_event_window_sizes_by_layer_group(self) -> Dict[int, int]: # mixed windows in one group, this needs to fan out per-layer. def get_event_window_size(layer_id: int) -> int: - window_size = self.kv_cache_manager_py_config.layers[layer_id].sliding_window_size + layer_config = self.kv_cache_manager_py_config.layers[layer_id] + window_size = getattr(layer_config, "sliding_window_size", None) return self.max_seq_len if window_size is None else int(window_size) return { @@ -1406,9 +1429,12 @@ def _prepare_swa_scratch_copy_tensors(self, index_mapper_capacity: int) -> None: for local_layer_idx in range(self.num_local_layers): layer_id = LayerId(local_layer_idx) pool_id = self.layer_to_pool_mapping_dict[layer_id] - roles = [Role.KEY, Role.VALUE] - if self.kv_cache_type == CacheTypeCpp.SELFKONLY: - roles[1] = Role.KEY + page_index_role = self._get_pool_page_index_role(pool_id) + paired_role = self._get_pool_paired_role(pool_id) + roles = [ + page_index_role, + page_index_role if paired_role is None else paired_role, + ] for role_idx, role in enumerate(roles): converter = self.impl.get_page_index_converter(layer_id, role) if converter.expansion != 1: @@ -3136,19 +3162,25 @@ def calculate_scaling_factor_size_bytes( ) return get_size_in_bytes(cache_size // quant_vector_size, scaling_factor_dtype) - def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool: - some_checks_unavailable = False - has_invalid_values = torch.tensor( - [False], dtype=torch.bool, device=torch.cuda.current_device() - ) + def _iter_cache_buffers_for_invalid_check(self) -> Iterable[torch.Tensor]: pool_handled = set() - - # Handle each layer from start to end to traverse the whole KV cache. for layer_id, layer_offset in self.layer_offsets.items(): pool_id = self.layer_to_pool_mapping_dict[layer_offset] if pool_id in pool_handled: continue buffer = self.get_buffers(layer_id) + if buffer is None: + continue + yield buffer + pool_handled.add(pool_id) + + def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool: + some_checks_unavailable = False + has_invalid_values = torch.tensor( + [False], dtype=torch.bool, device=torch.cuda.current_device() + ) + + for buffer in self._iter_cache_buffers_for_invalid_check(): # process in chunks of 256 pages to avoid OoM for i in range(0, buffer.shape[0], 256): buffer_slice = buffer[i : i + 256] @@ -3159,7 +3191,6 @@ def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool some_checks_unavailable = True if fill_with_zero: buffer.zero_() - pool_handled.add(pool_id) torch.cuda.synchronize() if some_checks_unavailable: diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 5357f709a649..96550e5fe53e 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -17,7 +17,8 @@ import os from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Tuple, Union +from typing import (TYPE_CHECKING, Dict, Iterable, List, NamedTuple, Optional, + Tuple, Union) import torch import triton @@ -52,8 +53,7 @@ KVCacheManagerConfig as KVCacheManagerConfigPy from tensorrt_llm.runtime.kv_cache_manager_v2 import (LayerId, PageIndexMode, SsmLayerConfig, - SwaScratchReuseConfig, - exact_div) + SwaScratchReuseConfig) GB = 1 << 30 @@ -225,6 +225,40 @@ class ReplayStateUpdateMetadata(NamedTuple): replay_history_size: int +def _advance_replay_state( + replay_metadata: ReplayStateUpdateMetadata, + state_indices: torch.Tensor, + accepted_tokens: torch.Tensor, + is_dummy_request: Optional[torch.Tensor] = None, +) -> None: + """Advance replay bookkeeping after a speculative generation step.""" + slots = state_indices.long() + accepted_tokens = accepted_tokens.to( + replay_metadata.prev_num_accepted_tokens.dtype) + prev_num_accepted_tokens = replay_metadata.prev_num_accepted_tokens[slots] + wrote_checkpoint = (prev_num_accepted_tokens + + replay_metadata.replay_step_width + > replay_metadata.replay_history_size) + next_num_accepted_tokens = torch.where( + wrote_checkpoint, + accepted_tokens, + prev_num_accepted_tokens + accepted_tokens, + ) + cache_buf_idx = replay_metadata.cache_buf_idx[slots] + next_cache_buf_idx = torch.where(wrote_checkpoint, 1 - cache_buf_idx, + cache_buf_idx) + if is_dummy_request is not None: + next_num_accepted_tokens = torch.where( + is_dummy_request, + prev_num_accepted_tokens, + next_num_accepted_tokens, + ) + next_cache_buf_idx = torch.where(is_dummy_request, cache_buf_idx, + next_cache_buf_idx) + replay_metadata.prev_num_accepted_tokens[slots] = next_num_accepted_tokens + replay_metadata.cache_buf_idx[slots] = next_cache_buf_idx + + class BaseMambaCacheManager(ABC): """Abstract interface for accessing mamba/recurrent state caches.""" @@ -860,33 +894,16 @@ def update_mamba_states(self, attn_metadata: "AttentionMetadata", src_state_indices = self.intermediate_state_indices[:num_gens] if self._use_replay_state_update: - # SSM state is handled incrementally by the kernel. Mirror the - # kernel's per-slot checkpoint predicate from the previous PNAT and - # fixed replay step width: checkpoint steps write a fresh history - # buffer and flip, while no-checkpoint steps append to the active - # buffer and keep reading from it next step. - accepted_tokens = num_accepted_tokens[num_contexts:num_contexts + - num_gens] - prev_num_accepted_tokens = \ - self.mamba_cache.prev_num_accepted_tokens[state_indices_d] - wrote_checkpoint = (prev_num_accepted_tokens + - self.replay_step_width - > self.replay_history_size) - next_num_accepted_tokens = torch.where( - wrote_checkpoint, accepted_tokens, - prev_num_accepted_tokens + accepted_tokens) - cache_buf_idx = self.mamba_cache.cache_buf_idx[state_indices_d] is_dummy_request = self._dummy_request_mask[ num_contexts:num_contexts + num_gens] - next_num_accepted_tokens = torch.where(is_dummy_request, - prev_num_accepted_tokens, - next_num_accepted_tokens) - self.mamba_cache.prev_num_accepted_tokens[state_indices_d] = \ - next_num_accepted_tokens - self.mamba_cache.cache_buf_idx[state_indices_d] = \ - torch.where(is_dummy_request, cache_buf_idx, - torch.where(wrote_checkpoint, 1 - cache_buf_idx, - cache_buf_idx)) + replay_metadata = self.get_replay_state_update_metadata() + assert replay_metadata is not None + _advance_replay_state( + replay_metadata, + state_indices_d, + num_accepted_tokens[num_contexts:num_contexts + num_gens], + is_dummy_request, + ) else: # Legacy: copy accepted SSM state from intermediate cache. ssm_states = self.mamba_cache.temporal @@ -1056,14 +1073,241 @@ def update_mamba_states(self, attn_metadata: "AttentionMetadata", class MambaHybridCacheManager(BaseResourceManager, BaseMambaCacheManager): - """Marker base class for hybrid mamba cache manager implementations. + """Shared interface and state plumbing for hybrid Mamba managers. - Used purely for ``isinstance`` / type-hint purposes so callers can refer - to the family without caring about the concrete implementation. Concrete - selection (V2, Mixed, or Cpp) lives in - ``_util.py:get_kv_cache_manager_cls``. + Concrete storage, state-index, and resource lifecycles remain owned by the + Cpp and V2 implementations. """ + _supports_additional_snapshot_offsets = False + + def _setup_mtp_intermediate_states(self, spec_config, + max_batch_size: int) -> None: + self.spec_config = spec_config + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None + if spec_config is None or self.local_num_mamba_layers == 0: + return + + tokens_per_gen_step = spec_config.tokens_per_gen_step + if not self._use_replay_state_update: + self.intermediate_ssm_states = torch.zeros( + size=[ + self.local_num_mamba_layers, max_batch_size, + tokens_per_gen_step + ] + self.ssm_state_shape, + dtype=self.ssm_state_dtype, + device="cuda", + ) + self.intermediate_conv_states = torch.zeros( + size=[ + self.local_num_mamba_layers, max_batch_size, tokens_per_gen_step + ] + self.conv_state_shape, + dtype=self.conv_state_dtype, + device="cuda", + ) + self.intermediate_state_indices = torch.arange(max_batch_size, + dtype=torch.int32, + device="cuda") + + def _allocate_pool_replay_buffers( + self, + spec_config, + cache_size: int, + device: Optional[torch.device], + ) -> bool: + """Allocate replay tensors shared by the Cpp and V2 state pools.""" + self.prev_num_accepted_tokens = None + self.cache_buf_idx = None + self.mamba_ssm_rand_seed = None + self.old_x = None + self.old_B = None + self.old_dt = None + self.old_dA_cumsum = None + + if (self.local_num_mamba_layers == 0 + or (not self._use_replay_state_update + and not self._mamba_ssm_stochastic_rounding)): + return False + + assert device is not None + self.mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( + cache_size, self._seed_rank_offset, device) + if spec_config is None or not self._use_replay_state_update: + return False + + history_size = self.replay_history_size + assert history_size is not None + nheads, head_dim, d_state = self.ssm_state_shape + common_shape = [self.local_num_mamba_layers, cache_size, 2] + self.prev_num_accepted_tokens = torch.zeros(cache_size, + dtype=torch.int32, + device=device) + self.cache_buf_idx = torch.zeros(cache_size, + dtype=torch.int32, + device=device) + self.old_x = torch.zeros( + common_shape + [history_size, nheads, head_dim], + dtype=self.conv_state_dtype, + device=device, + ) + self.old_B = torch.zeros( + common_shape + [history_size, self._n_groups_per_rank, d_state], + dtype=self.conv_state_dtype, + device=device, + ) + self.old_dt = torch.zeros( + common_shape + [nheads, history_size], + dtype=torch.float32, + device=device, + ) + self.old_dA_cumsum = torch.zeros( + common_shape + [nheads, history_size], + dtype=torch.float32, + device=device, + ) + return True + + def _reset_context_mamba_slots(self, num_contexts: int) -> None: + if num_contexts == 0: + return + + context_slots = self.cuda_state_indices[:num_contexts].long() + if (self._use_replay_state_update + and self.prev_num_accepted_tokens is not None + and self.cache_buf_idx is not None): + self.prev_num_accepted_tokens[context_slots] = 0 + self.cache_buf_idx[context_slots] = 0 + if self.old_x is not None: + self.old_x[:, context_slots] = 0 + if self.old_B is not None: + self.old_B[:, context_slots] = 0 + if self.old_dt is not None: + self.old_dt[:, context_slots] = 0 + if self.old_dA_cumsum is not None: + self.old_dA_cumsum[:, context_slots] = 0 + + if self.mamba_ssm_rand_seed is None: + return + self._seed_request_counter += 1 + counter = self._seed_request_counter + rank_offset = self._seed_rank_offset + host_slots = self._host_state_indices[:num_contexts].tolist() + new_seeds = [ + _compute_deterministic_mamba_seed(counter, slot, rank_offset) + for slot in host_slots + ] + self.mamba_ssm_rand_seed[context_slots] = torch.tensor( + new_seeds, + dtype=torch.int64, + device=self.mamba_ssm_rand_seed.device, + ) + + def prepare_expect_snapshot_points(self, + requests: List[LlmRequest]) -> None: + """Set reusable Mamba snapshot boundaries before scheduling.""" + if not self.kv_cache_config.enable_block_reuse: + for request in requests: + request.expect_snapshot_points = [] + return + + state_config = self.kv_cache_config.mamba_state_config + interval = state_config.periodic_snapshot_interval + for request in requests: + snapshot_points = set() + if interval is not None and interval > 0: + snapshot_points.update( + range(interval, request.prompt_len + 1, interval)) + if self._supports_additional_snapshot_offsets: + for offset in state_config.additional_snapshot_offsets_from_start: + if offset <= request.prompt_len: + snapshot_points.add(offset) + for offset in state_config.additional_snapshot_offsets_from_end: + point = request.prompt_len - offset + if point > 0: + snapshot_points.add(point) + request.expect_snapshot_points = sorted(snapshot_points) + + def is_speculative(self) -> bool: + return self.spec_config is not None + + def get_ssm_states(self, layer_idx: int) -> torch.Tensor: + return self.all_ssm_states[self.mamba_layer_offsets[layer_idx]] + + def get_conv_states(self, layer_idx: int) -> torch.Tensor: + return self.all_conv_states[self.mamba_layer_offsets[layer_idx]] + + def get_intermediate_ssm_states(self, + layer_idx: int) -> Optional[torch.Tensor]: + if self.intermediate_ssm_states is None: + return None + return self.intermediate_ssm_states[self.mamba_layer_offsets[layer_idx]] + + def get_intermediate_conv_states(self, + layer_idx: int) -> Optional[torch.Tensor]: + if self.intermediate_conv_states is None: + return None + return self.intermediate_conv_states[ + self.mamba_layer_offsets[layer_idx]] + + def mamba_layer_cache( + self, layer_idx: int + ) -> Union[PythonMambaCacheManager.State, + PythonMambaCacheManager.SpeculativeState, None]: + conv = self.get_conv_states(layer_idx) + ssm = self.get_ssm_states(layer_idx) + if self.spec_config is not None: + layer_offset = self.mamba_layer_offsets[layer_idx] + spec_kwargs = {} + if self.mamba_ssm_rand_seed is not None: + spec_kwargs['mamba_ssm_rand_seed'] = self.mamba_ssm_rand_seed + if self._use_replay_state_update: + spec_kwargs['old_x'] = self.old_x[layer_offset] + spec_kwargs['old_B'] = self.old_B[layer_offset] + spec_kwargs['old_dt'] = self.old_dt[layer_offset] + spec_kwargs['old_dA_cumsum'] = self.old_dA_cumsum[layer_offset] + spec_kwargs['cache_buf_idx'] = self.cache_buf_idx + spec_kwargs['prev_num_accepted_tokens'] = ( + self.prev_num_accepted_tokens) + else: + spec_kwargs['intermediate_ssm'] = self.intermediate_ssm_states[ + layer_offset] + return PythonMambaCacheManager.SpeculativeState( + conv=conv, + temporal=ssm, + intermediate_conv_window=self. + intermediate_conv_states[layer_offset], + **spec_kwargs, + ) + return PythonMambaCacheManager.State(conv=conv, temporal=ssm) + + @property + def use_replay_state_update(self) -> bool: + return self.get_replay_state_update_metadata() is not None + + def get_replay_state_update_metadata( + self) -> Optional[ReplayStateUpdateMetadata]: + prev_num_accepted_tokens = getattr(self, 'prev_num_accepted_tokens', + None) + cache_buf_idx = getattr(self, 'cache_buf_idx', None) + if (not self._use_replay_state_update + or prev_num_accepted_tokens is None or cache_buf_idx is None + or self.replay_step_width is None + or self.replay_history_size is None): + return None + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=prev_num_accepted_tokens, + cache_buf_idx=cache_buf_idx, + replay_step_width=self.replay_step_width, + replay_history_size=self.replay_history_size) + + def get_mamba_ssm_cache_dtype(self) -> torch.dtype: + return self.ssm_state_dtype + + def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: + return getattr(self, 'mamba_ssm_rand_seed', None) + def _get_mamba_hybrid_pool_size(max_batch_size: int, mapping: Mapping) -> int: """Return the internal Mamba state pool size for MixedMambaHybridCacheManager.""" @@ -1545,6 +1789,12 @@ def __init__( # seed values without any torch.randint. self._seed_request_counter = 0 self.ssm_state_dtype = mamba_ssm_cache_dtype + # Keep the shared Mamba interface valid on PP ranks that do not own a + # local Mamba layer. + self.spec_config = spec_config + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None if self.local_num_mamba_layers == 0: logger.info( @@ -1915,47 +2165,7 @@ def _prepare_resources(self, scheduled_batch: ScheduledRequests): self._pending_state_transfers = self.impl.copy_linear_attention_block_batch( self.requests) self._setup_state_indices() - # Reset replay double-buffer state for fresh context blocks. A reused - # block (prefix-cache hit or block recycled across requests) may carry - # stale prev_num_accepted_tokens / cache_buf_idx values from a prior - # owner; the replay kernel reads these on the first decode step. - num_contexts = len(scheduled_batch.context_requests) - if num_contexts > 0: - ctx_slots = self.cuda_state_indices[:num_contexts].long() - if (self._use_replay_state_update - and self.prev_num_accepted_tokens is not None - and self.cache_buf_idx is not None): - self.prev_num_accepted_tokens[ctx_slots] = 0 - self.cache_buf_idx[ctx_slots] = 0 - if self.old_x is not None: - self.old_x[:, ctx_slots] = 0 - if self.old_B is not None: - self.old_B[:, ctx_slots] = 0 - if self.old_dt is not None: - self.old_dt[:, ctx_slots] = 0 - if self.old_dA_cumsum is not None: - self.old_dA_cumsum[:, ctx_slots] = 0 - # Deterministic per-context-slot seed rotation. Runs whenever - # the seed buffer exists, including the non-replay SR path. - # Bump the host counter once per batch and write one new seed - # per fresh context slot from a pure function of - # (counter, slot, rank). No torch.randint involved. - if self.mamba_ssm_rand_seed is not None: - self._seed_request_counter += 1 - counter = self._seed_request_counter - rank_offset = self._seed_rank_offset - host_slots = ctx_slots.cpu().tolist() - new_seeds = [ - _compute_deterministic_mamba_seed(counter, slot, - rank_offset) - for slot in host_slots - ] - seed_tensor = torch.tensor( - new_seeds, - dtype=torch.int64, - device=self.mamba_ssm_rand_seed.device, - ) - self.mamba_ssm_rand_seed[ctx_slots] = seed_tensor + self._reset_context_mamba_slots(len(scheduled_batch.context_requests)) def prepare_resources(self, scheduled_batch: ScheduledRequests): super().prepare_resources(scheduled_batch) @@ -1973,9 +2183,6 @@ def flush_state_transfers(self) -> None: self.impl.refresh_blocks() self._pending_state_transfers = False - def is_speculative(self) -> bool: - return self.spec_config is not None - @nvtx_range("hybrid_update_mamba_states") def update_mamba_states(self, attn_metadata: "AttentionMetadata", @@ -2011,32 +2218,17 @@ def update_mamba_states(self, # in one launch before PNAT and the active history buffer change. self._commit_gdn_cached_replay_history_layers( attn_metadata, num_gens) - # SSM state is handled incrementally by the kernel. Mirror the - # kernel's checkpoint predicate from the previous PNAT and fixed - # replay step width: checkpoint steps flip buffers, while no-write - # steps append to the active history. - slots = state_indices_d.long() - accepted = num_accepted_tokens[num_contexts:num_contexts + - num_gens].to( - self.prev_num_accepted_tokens. - dtype) - prev_num_accepted_tokens = self.prev_num_accepted_tokens[slots] - wrote_checkpoint = (prev_num_accepted_tokens + - self.replay_step_width - > self.replay_history_size) - next_num_accepted_tokens = torch.where( - wrote_checkpoint, accepted, prev_num_accepted_tokens + accepted) - cache_buf_idx = self.cache_buf_idx[slots] assert self._dummy_request_mask is not None is_dummy_request = self._dummy_request_mask[ num_contexts:num_contexts + num_gens] - next_num_accepted_tokens = torch.where(is_dummy_request, - prev_num_accepted_tokens, - next_num_accepted_tokens) - self.prev_num_accepted_tokens[slots] = next_num_accepted_tokens - self.cache_buf_idx[slots] = torch.where( - is_dummy_request, cache_buf_idx, - torch.where(wrote_checkpoint, 1 - cache_buf_idx, cache_buf_idx)) + replay_metadata = self.get_replay_state_update_metadata() + assert replay_metadata is not None + _advance_replay_state( + replay_metadata, + state_indices_d, + num_accepted_tokens[num_contexts:num_contexts + num_gens], + is_dummy_request, + ) else: # Legacy: copy the accepted SSM state from the intermediate buffer. _promote_mamba_state_triton(self.all_ssm_states, @@ -2096,64 +2288,6 @@ def get_num_available_tokens(self, result = min(result, rs_token_cap) return max(result, 0) - def get_ssm_states(self, layer_idx: int) -> torch.Tensor: - return self.all_ssm_states[self.mamba_layer_offsets[layer_idx]] - - def get_conv_states(self, layer_idx: int) -> torch.Tensor: - return self.all_conv_states[self.mamba_layer_offsets[layer_idx]] - - def get_intermediate_ssm_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - if self.intermediate_ssm_states is None: - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.intermediate_ssm_states[layer_offset] - - def get_intermediate_conv_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - if self.intermediate_conv_states is None: - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.intermediate_conv_states[layer_offset] - - def mamba_layer_cache( - self, layer_idx: int - ) -> Union[PythonMambaCacheManager.State, - PythonMambaCacheManager.SpeculativeState, None]: - conv = self.get_conv_states(layer_idx) - ssm = self.get_ssm_states(layer_idx) - if self.spec_config is not None: - layer_offset = self.mamba_layer_offsets[layer_idx] - spec_kwargs = {} - # Per-cache-slot Philox seed buffer is shared across replay and - # non-replay MTP paths. The mixer asserts non-None on both - # branches when SR is enabled, so pass it through whenever it - # exists — not just on the replay branch. - if self.mamba_ssm_rand_seed is not None: - spec_kwargs['mamba_ssm_rand_seed'] = self.mamba_ssm_rand_seed - if self._use_replay_state_update: - # Per-layer slices for the replay kernel; shared 1D tensors - # (cache_buf_idx, prev_num_accepted_tokens) are passed - # untouched via the SpeculativeState._SHARED_FIELDS contract. - spec_kwargs['old_x'] = self.old_x[layer_offset] - spec_kwargs['old_B'] = self.old_B[layer_offset] - spec_kwargs['old_dt'] = self.old_dt[layer_offset] - spec_kwargs['old_dA_cumsum'] = self.old_dA_cumsum[layer_offset] - spec_kwargs['cache_buf_idx'] = self.cache_buf_idx - spec_kwargs['prev_num_accepted_tokens'] = ( - self.prev_num_accepted_tokens) - else: - spec_kwargs['intermediate_ssm'] = self.intermediate_ssm_states[ - layer_offset] - return PythonMambaCacheManager.SpeculativeState( - conv=conv, - temporal=ssm, - intermediate_conv_window=self. - intermediate_conv_states[layer_offset], - **spec_kwargs, - ) - return PythonMambaCacheManager.State(conv=conv, temporal=ssm) - def free_resources(self, request: LlmRequest, pin_on_release: bool = False): if request in self.requests: self.requests.remove(request) @@ -2254,24 +2388,6 @@ def get_state_indices(self, return indices return self.cuda_state_indices - def prepare_expect_snapshot_points(self, - requests: List[LlmRequest]) -> None: - """Set reusable Mamba snapshot boundaries before scheduling.""" - if not self.kv_cache_config.enable_block_reuse: - for request in requests: - request.expect_snapshot_points = [] - return - - interval = self.linear_attention_metadata.states_snapshot_interval - if interval is None or interval <= 0: - for request in requests: - request.expect_snapshot_points = [] - return - - for request in requests: - request.expect_snapshot_points = list( - range(interval, request.prompt_len + 1, interval)) - def _setup_states(self) -> None: # Pool layout: {numLocalLayers, numBlocks, ssm_bytes + conv_bytes} (as uint8) pool: torch.Tensor = self.impl.get_recurrent_states_pool().view( @@ -2291,167 +2407,23 @@ def _setup_states(self) -> None: self.all_ssm_states.zero_() self.all_conv_states.zero_() - def _setup_mtp_intermediate_states(self, spec_config, - max_batch_size) -> None: - self.spec_config = spec_config - self.intermediate_ssm_states = None - self.intermediate_conv_states = None - self.intermediate_state_indices = None - if self.spec_config is not None: - # DFlash/PARD use 2K query tokens per gen, so size by tokens_per_gen_step. - speculative_num_draft_tokens = self.spec_config.tokens_per_gen_step - 1 - num_local_mamba_layers = len(self.mamba_pp_layers) - - # Legacy SSM intermediate buffer is only needed when replay is - # disabled; replay reads from the per-block double-buffered cache - # set up in _setup_replay_buffers instead. - if not self._use_replay_state_update: - self.intermediate_ssm_states = torch.zeros( - size=[ - num_local_mamba_layers, max_batch_size, - speculative_num_draft_tokens + 1 - ] + self.ssm_state_shape, - dtype=self.ssm_state_dtype, - device="cuda", - ) - - self.intermediate_conv_states = torch.zeros( - size=[ - num_local_mamba_layers, max_batch_size, - speculative_num_draft_tokens + 1 - ] + self.conv_state_shape, - dtype=self.conv_state_dtype, - device="cuda", - ) - - self.intermediate_state_indices = torch.arange(max_batch_size, - dtype=torch.int32, - device="cuda") - def _setup_replay_buffers(self, spec_config) -> None: - """Allocate per-pool-block replay buffers used by replay_selective_state_update. - - Unlike the Mixed cache manager (where slots are 0..max_batch_size-1), - the unified C++ KV pool assigns recurrent-state block indices up to - ``num_blocks_in_pool``. The replay kernel indexes ``cache_buf_idx`` and - ``prev_num_accepted_tokens`` by these block indices, so the buffers - must match the pool extent rather than ``max_batch_size``. - """ - # Replay tensors require spec_config + replay path enabled. The - # rand_seed buffer is separable from replay and must also be - # allocated for non-replay SR so the flashinfer path has a - # persistent deterministic seed source. - self.prev_num_accepted_tokens = None - self.cache_buf_idx = None - self.mamba_ssm_rand_seed = None - self.old_x = None - self.old_B = None - self.old_dt = None - self.old_dA_cumsum = None - - if (not self._use_replay_state_update - and not self._mamba_ssm_stochastic_rounding): - return - cache_size = self.all_ssm_states.shape[1] device = self.all_ssm_states.device - # Always-available deterministic seed buffer when SR (or replay) - # is on. Works for non-MTP runs because we don't depend on - # spec_config to allocate it. - self.mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( - cache_size, self._seed_rank_offset, device) - - if spec_config is None or not self._use_replay_state_update: - # Without spec_config or replay we still keep the seed buffer - # (above) so the non-MTP flashinfer SR path has a persistent - # rand_seed source. - self.prev_num_accepted_tokens = None - self.cache_buf_idx = None - self.old_x = None - self.old_B = None - self.old_dt = None - self.old_dA_cumsum = None - self._dummy_request_mask = None - self._dummy_request_mask_host = None + self._dummy_request_mask = None + self._dummy_request_mask_host = None + if not self._allocate_pool_replay_buffers(spec_config, cache_size, + device): return - history_size = self.replay_history_size - num_local_mamba_layers = self.local_num_mamba_layers - nheads, head_dim, d_state = self.ssm_state_shape - n_groups_per_rank = self._n_groups_per_rank - - # Shared across layers (consumed by the replay kernel via slot index). - self.prev_num_accepted_tokens = torch.zeros(cache_size, - dtype=torch.int32, - device=device) - self.cache_buf_idx = torch.zeros(cache_size, - dtype=torch.int32, - device=device) self._dummy_request_mask = torch.zeros(self.max_batch_size, dtype=torch.bool, device=device) - self._dummy_request_mask_host = torch.zeros(self.max_batch_size, - dtype=torch.bool, - pin_memory=prefer_pinned()) - self.old_x = torch.zeros(num_local_mamba_layers, - cache_size, - 2, - history_size, - nheads, - head_dim, - dtype=self.conv_state_dtype, - device=device) - # Per-layer double-buffered caches. - self.old_B = torch.zeros(num_local_mamba_layers, - cache_size, - 2, - history_size, - n_groups_per_rank, - d_state, - dtype=self.conv_state_dtype, - device=device) - self.old_dt = torch.zeros(num_local_mamba_layers, - cache_size, - 2, - nheads, - history_size, - dtype=torch.float32, - device=device) - self.old_dA_cumsum = torch.zeros(num_local_mamba_layers, - cache_size, - 2, - nheads, - history_size, - dtype=torch.float32, - device=device) - - @property - def use_replay_state_update(self) -> bool: - return self.get_replay_state_update_metadata() is not None - - def get_replay_state_update_metadata( - self) -> Optional[ReplayStateUpdateMetadata]: - prev_num_accepted_tokens = getattr(self, 'prev_num_accepted_tokens', - None) - cache_buf_idx = getattr(self, 'cache_buf_idx', None) - if (not self._use_replay_state_update - or prev_num_accepted_tokens is None or cache_buf_idx is None - or self.replay_step_width is None - or self.replay_history_size is None): - return None - return ReplayStateUpdateMetadata( - prev_num_accepted_tokens=prev_num_accepted_tokens, - cache_buf_idx=cache_buf_idx, - replay_step_width=self.replay_step_width, - replay_history_size=self.replay_history_size) - - def get_mamba_ssm_cache_dtype(self) -> torch.dtype: - return self.ssm_state_dtype - - def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: - """Return the persistent (cache_size,) int64 Philox seed buffer or - None when stochastic rounding is not active for this manager.""" - return getattr(self, 'mamba_ssm_rand_seed', None) + self._dummy_request_mask_host = torch.zeros( + self.max_batch_size, + dtype=torch.bool, + pin_memory=prefer_pinned(), + ) class V2MambaHybridCacheManager(KVCacheManagerV2, MambaHybridCacheManager): @@ -2463,6 +2435,8 @@ class V2MambaHybridCacheManager(KVCacheManagerV2, MambaHybridCacheManager): the PyTorch Mamba kernels. """ + _supports_additional_snapshot_offsets = True + def __init__( self, # mamba cache parameters @@ -2510,7 +2484,6 @@ def __init__( ] self._mamba_layer_mask = list(mamba_layer_mask) - self.requests = [] self._use_replay_state_update = use_replay_state_update self.replay_step_width: Optional[int] = ( spec_config.tokens_per_gen_step @@ -2698,291 +2671,7 @@ def _get_pool_paired_role(self, pool_id: int) -> Optional[DataRole]: layer_id = int(self.impl.layer_grouping[pool_id][0]) if self._is_local_mamba_layer(layer_id): return None - return (Role.VALUE - if self.kv_cache_type != CacheTypeCpp.SELFKONLY else None) - - def _get_block_scale_role_for_pool(self, - pool_id: int) -> Optional[DataRole]: - if (self.dtype != DataType.NVFP4 - or self._get_pool_page_index_role(pool_id) != Role.KEY): - return None - return Role.KEY_BLOCK_SCALE - - def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: - """Build attention-compatible tables for mixed attention/state pools.""" - kv_cache_pool_pointers_list = [] - kv_cache_pool_mapping_list = [] - block_scale_pool_pointers_list = [] - if self.enable_swa_scratch_reuse: - for local_layer_idx in range(self.num_local_layers): - layer_id = LayerId(local_layer_idx) - pool_id = self.impl.get_layer_group_id(layer_id) - page_index_role = self._get_pool_page_index_role(pool_id) - kv_cache_pool_pointers_list.append([ - self.impl.get_mem_pool_base_address( - layer_id, page_index_role, PageIndexMode.PER_LAYER), - 0, - ]) - if self.dtype == DataType.NVFP4: - block_scale_role = self._get_block_scale_role_for_pool( - pool_id) - block_scale_pool_pointers_list.append([ - self.impl.get_mem_pool_base_address( - layer_id, block_scale_role, PageIndexMode.PER_LAYER) - if block_scale_role is not None else 0, - 0, - ]) - kv_cache_pool_mapping_list.append([int(layer_id), 0]) - else: - for pool_id in range(self.num_pools): - layer_id = self.impl.layer_grouping[pool_id][0] - page_index_role = self._get_pool_page_index_role(pool_id) - kv_cache_pool_pointers_list.append([ - self.impl.get_mem_pool_base_address(layer_id, - page_index_role, - PageIndexMode.SHARED), - 0, - ]) - if self.dtype == DataType.NVFP4: - block_scale_role = self._get_block_scale_role_for_pool( - pool_id) - block_scale_pool_pointers_list.append([ - self.impl.get_mem_pool_base_address( - layer_id, block_scale_role, PageIndexMode.SHARED) - if block_scale_role is not None else 0, - 0, - ]) - - for local_layer_idx in range(self.num_local_layers): - layer_id = LayerId(local_layer_idx) - layer_group_id = self.impl.get_layer_group_id(layer_id) - page_index_role = self._get_pool_page_index_role(layer_group_id) - index_base_addr = kv_cache_pool_pointers_list[layer_group_id][0] - addr_offset = (self.impl.get_mem_pool_base_address( - layer_id, page_index_role, PageIndexMode.SHARED) - - index_base_addr) - offset_divisor = self.impl.get_page_stride( - layer_id, page_index_role) - if self._get_pool_paired_role(layer_group_id) is not None: - offset_divisor *= self.kv_factor - offset = exact_div(addr_offset, offset_divisor) - - if self.dtype != DataType.NVFP4: - block_scale_offset = None - else: - block_scale_role = self._get_block_scale_role_for_pool( - layer_group_id) - if block_scale_role is None: - block_scale_offset = None - else: - block_scale_base_addr = block_scale_pool_pointers_list[ - layer_group_id][0] - block_scale_addr_offset = ( - self.impl.get_mem_pool_base_address( - layer_id, block_scale_role, - PageIndexMode.SHARED) - block_scale_base_addr) - block_scale_divisor = self.impl.get_page_stride( - layer_id, block_scale_role) - if self._get_pool_paired_role( - layer_group_id) is not None: - block_scale_divisor *= self.kv_factor - block_scale_offset = exact_div(block_scale_addr_offset, - block_scale_divisor) - - if block_scale_offset is not None: - assert block_scale_offset == offset, ( - "Block scale offset and offset should be the same") - - kv_cache_pool_mapping_list.append([layer_group_id, offset]) - - if self.dtype == DataType.NVFP4: - for pool_id, block_scale_pool_pointers in enumerate( - block_scale_pool_pointers_list): - pool_pointers = kv_cache_pool_pointers_list[pool_id] - kv_cache_pool_pointers_list[pool_id] = [ - [pool_pointers[0], block_scale_pool_pointers[0]], - [pool_pointers[1], block_scale_pool_pointers[1]], - ] - - self.kv_cache_pool_pointers = torch.tensor( - kv_cache_pool_pointers_list, - dtype=torch.int64, - device="cpu", - pin_memory=prefer_pinned(), - ) - self.kv_cache_pool_mapping = torch.tensor( - kv_cache_pool_mapping_list, - dtype=torch.int32, - device="cpu", - pin_memory=prefer_pinned(), - ) - self.index_scales = torch.empty( - self.num_pools, - dtype=torch.int32, - pin_memory=prefer_pinned(), - device="cpu", - ) - self.kv_offset = torch.empty( - self.num_pools, - dtype=torch.int32, - pin_memory=prefer_pinned(), - device="cpu", - ) - for pool_id in range(self.num_pools): - layer_id = self.impl.layer_grouping[pool_id][0] - page_index_role = self._get_pool_page_index_role(pool_id) - self.index_scales[pool_id] = self.impl.get_page_index_scale( - layer_id, page_index_role) - paired_role = self._get_pool_paired_role(pool_id) - if paired_role is not None: - self.kv_offset[pool_id] = exact_div( - self.impl.get_mem_pool_base_address(layer_id, paired_role, - PageIndexMode.SHARED) - - self.impl.get_mem_pool_base_address( - layer_id, page_index_role, PageIndexMode.SHARED), - self.impl.get_page_stride(layer_id, page_index_role), - ) - else: - self.kv_offset[pool_id] = 0 - self._index_scale_ints = self.index_scales.tolist() - - self.host_kv_cache_block_offsets = torch.zeros( - self.num_pools, - index_mapper_capacity * self.max_beam_width, - 2, - self.max_blocks_per_seq, - dtype=torch.int32, - pin_memory=prefer_pinned(), - device="cpu", - ) - if self.enable_swa_scratch_reuse: - self._prepare_swa_scratch_copy_tensors(index_mapper_capacity) - - def _prepare_swa_scratch_copy_tensors(self, - index_mapper_capacity: int) -> None: - pool_ids = torch.empty( - self.num_attention_op_pools, - 2, - dtype=torch.long, - device="cpu", - ) - scales = torch.empty( - self.num_attention_op_pools, - 2, - dtype=torch.int32, - device="cpu", - ) - layer_offsets = torch.empty_like(scales) - scratch_pages = torch.empty_like(scales) - for local_layer_idx in range(self.num_local_layers): - layer_id = LayerId(local_layer_idx) - pool_id = self.layer_to_pool_mapping_dict[layer_id] - page_index_role = self._get_pool_page_index_role(pool_id) - paired_role = self._get_pool_paired_role(pool_id) - roles = [page_index_role, paired_role or page_index_role] - for role_idx, role in enumerate(roles): - converter = self.impl.get_page_index_converter(layer_id, role) - if converter.expansion != 1: - raise NotImplementedError( - "SWA scratch block-table conversion does not support " - "expanded page indices yet: " - f"layer={layer_id}, role={role}, " - f"expansion={converter.expansion}") - pool_ids[local_layer_idx, role_idx] = pool_id - scales[local_layer_idx, role_idx] = int(converter.scale) - layer_offsets[local_layer_idx, - role_idx] = int(converter.layer_offset) - scratch_pages[local_layer_idx, - role_idx] = int(converter.scratch_pages_per_block) - - staging_capacity = index_mapper_capacity * self.max_beam_width - device = torch.device("cuda", torch.cuda.current_device()) - self._device_kv_cache_block_offsets_input = torch.empty_like( - self.host_kv_cache_block_offsets, device=device) - self._device_attention_op_block_offsets_staging = torch.empty( - self.num_attention_op_pools, - staging_capacity, - 2, - self.max_blocks_per_seq, - dtype=torch.int32, - device=device, - ) - self._device_copy_idx_staging = torch.zeros(staging_capacity, - dtype=torch.long, - device=device) - self._device_num_contexts = torch.empty((), - dtype=torch.int32, - device=device) - self._device_attention_op_pool_ids = pool_ids.to(device=device) - self._device_attention_op_scales = scales.to(device=device) - self._device_attention_op_layer_offsets = layer_offsets.to( - device=device) - self._device_attention_op_scratch_pages = scratch_pages.to( - device=device) - self._device_block_positions = torch.arange(self.max_blocks_per_seq, - dtype=torch.int32, - device=device) - - min_scale = int(scales.min().item()) if scales.numel() > 0 else 1 - max_scratch_pages = (int(scratch_pages.max().item()) - if scratch_pages.numel() > 0 else 1) - self._max_scratch_slots = max( - 1, - (self.max_blocks_per_seq * max_scratch_pages + min_scale - 1) // - min_scale, - ) - scratch_slots_shape = (self.num_pools, staging_capacity, - self._max_scratch_slots) - self._host_scratch_begs_staging = torch.zeros( - self.num_pools, - staging_capacity, - dtype=torch.int32, - pin_memory=prefer_pinned(), - device="cpu", - ) - self._host_scratch_ends_staging = torch.zeros( - self.num_pools, - staging_capacity, - dtype=torch.int32, - pin_memory=prefer_pinned(), - device="cpu", - ) - self._host_scratch_slots_staging = torch.zeros( - scratch_slots_shape, - dtype=torch.int32, - pin_memory=prefer_pinned(), - device="cpu", - ) - self._device_scratch_begs_staging = torch.zeros( - self.num_pools, - staging_capacity, - dtype=torch.int32, - device=device, - ) - self._device_scratch_ends_staging = torch.zeros( - self.num_pools, - staging_capacity, - dtype=torch.int32, - device=device, - ) - self._device_scratch_slots_staging = torch.zeros( - scratch_slots_shape, - dtype=torch.int32, - device=device, - ) - - def _get_event_window_sizes_by_layer_group(self) -> Dict[int, int]: - - def get_event_window_size(layer_id: int) -> int: - layer_config = self.kv_cache_manager_py_config.layers[layer_id] - window_size = getattr(layer_config, "sliding_window_size", None) - return (self.max_seq_len - if window_size is None else int(window_size)) - - return { - int(layer_group_id): get_event_window_size(int(layer_ids[0])) - for layer_group_id, layer_ids in enumerate(self.impl.layer_grouping) - } + return super()._get_pool_paired_role(pool_id) def _max_resident_sequences(self) -> int: return self.max_batch_size * self.mapping.pp_size @@ -3013,30 +2702,25 @@ def _ssm_slots_per_request_for_typical_batch( kv_cache_config: KvCacheConfig, ) -> List[int]: num_sequences = self._max_resident_sequences() - fixed_rules = 0 - if capacity > 0: - fixed_rules, _ = _mamba_snapshot_rule_counts( - kv_cache_config, self.max_seq_len, self.tokens_per_block) - interval = _mamba_regular_snapshot_interval(kv_cache_config, - self.max_seq_len) - regular_snapshots = capacity // interval if interval is not None else 0 - regular_per_request, regular_remainder = divmod(regular_snapshots, - num_sequences) + snapshot_slots = self._num_ssm_snapshots_for_capacity( + capacity, kv_cache_config) + snapshot_slots_per_request, snapshot_remainder = divmod( + snapshot_slots, num_sequences) dummy_per_request, dummy_remainder = divmod( self._num_reserved_dummy_slots, num_sequences) return [ - 1 + fixed_rules + regular_per_request + int(i < regular_remainder) + + 1 + snapshot_slots_per_request + int(i < snapshot_remainder) + dummy_per_request + int(i < dummy_remainder) for i in range(num_sequences) ] def _get_quota_from_max_tokens(self, max_tokens: int) -> int: attention_quota = super()._get_quota_from_max_tokens(max_tokens) - ssm_slots = self._ssm_slots_per_request_for_typical_batch( - max_tokens, self.kv_cache_config) - state_slots = sum(ssm_slots) + num_request_lineages = self._max_resident_sequences() + state_slots = (num_request_lineages + self._num_reserved_dummy_slots + + self._num_ssm_snapshots_for_capacity( + max_tokens, self.kv_cache_config)) state_quota = state_slots * self._mamba_state_bytes_per_slot() - num_request_lineages = len(ssm_slots) # Once the plan contains any non-live SSM capacity, reserve one partial # attention page per request lineage. This remains conservative when # the plan contains fewer than one non-live slot per lineage. @@ -3234,96 +2918,14 @@ def _setup_states(self) -> None: for local_layer_idx in self.mamba_local_layer_ids ] - def _setup_mtp_intermediate_states(self, spec_config, - max_batch_size) -> None: - self.spec_config = spec_config - self.intermediate_ssm_states = None - self.intermediate_conv_states = None - self.intermediate_state_indices = None - if self.spec_config is not None and self.local_num_mamba_layers > 0: - speculative_num_draft_tokens = self.spec_config.tokens_per_gen_step - 1 - if not self._use_replay_state_update: - self.intermediate_ssm_states = torch.zeros( - size=[ - self.local_num_mamba_layers, max_batch_size, - speculative_num_draft_tokens + 1 - ] + self.ssm_state_shape, - dtype=self.ssm_state_dtype, - device="cuda", - ) - self.intermediate_conv_states = torch.zeros( - size=[ - self.local_num_mamba_layers, max_batch_size, - speculative_num_draft_tokens + 1 - ] + self.conv_state_shape, - dtype=self.conv_state_dtype, - device="cuda", - ) - self.intermediate_state_indices = torch.arange(max_batch_size, - dtype=torch.int32, - device="cuda") - def _setup_replay_buffers(self, spec_config) -> None: - self.prev_num_accepted_tokens = None - self.cache_buf_idx = None - self.mamba_ssm_rand_seed = None - self.old_x = None - self.old_B = None - self.old_dt = None - self.old_dA_cumsum = None - - if (self.local_num_mamba_layers == 0 - or (not self._use_replay_state_update - and not self._mamba_ssm_stochastic_rounding)): - return - - cache_size = self.all_ssm_states[0].shape[0] - assert all(t.shape[0] == cache_size for t in self.all_ssm_states) - device = self.all_ssm_states[0].device - self.mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( - cache_size, self._seed_rank_offset, device) - - if spec_config is None or not self._use_replay_state_update: - return - - T = spec_config.tokens_per_gen_step - nheads, head_dim, d_state = self.ssm_state_shape - self.prev_num_accepted_tokens = torch.zeros(cache_size, - dtype=torch.int32, - device=device) - self.cache_buf_idx = torch.zeros(cache_size, - dtype=torch.int32, - device=device) - self.old_x = torch.zeros(self.local_num_mamba_layers, - cache_size, - 2, - T, - nheads, - head_dim, - dtype=self.conv_state_dtype, - device=device) - self.old_B = torch.zeros(self.local_num_mamba_layers, - cache_size, - 2, - T, - self._n_groups_per_rank, - d_state, - dtype=self.conv_state_dtype, - device=device) - self.old_dt = torch.zeros(self.local_num_mamba_layers, - cache_size, - 2, - nheads, - T, - dtype=torch.float32, - device=device) - self.old_dA_cumsum = torch.zeros(self.local_num_mamba_layers, - cache_size, - 2, - nheads, - T, - dtype=torch.float32, - device=device) + cache_size = 0 + device = None + if self.local_num_mamba_layers > 0: + cache_size = self.all_ssm_states[0].shape[0] + assert all(t.shape[0] == cache_size for t in self.all_ssm_states) + device = self.all_ssm_states[0].device + self._allocate_pool_replay_buffers(spec_config, cache_size, device) def _attention_cache_bytes_per_token(self) -> int: # Mamba layers have zero KV heads, so the generic calculation naturally @@ -3384,47 +2986,17 @@ def get_buffers(self, return None return super().get_buffers(layer_idx, kv_layout) - def check_invalid_values_in_kv_cache(self, - fill_with_zero: bool = False) -> bool: - """Check both attention pages and recurrent-state storage.""" - some_checks_unavailable = False - has_invalid_values = torch.tensor([False], - dtype=torch.bool, - device=torch.cuda.current_device()) - - def check_buffer(buffer: torch.Tensor) -> None: - nonlocal some_checks_unavailable - for start in range(0, buffer.shape[0], 256): - buffer_slice = buffer[start:start + 256] - try: - has_invalid_values.logical_or_( - torch.isnan(buffer_slice).any()) - has_invalid_values.logical_or_( - torch.isinf(buffer_slice).any()) - except NotImplementedError: - some_checks_unavailable = True - if fill_with_zero: - buffer.zero_() - + def _iter_cache_buffers_for_invalid_check(self) -> Iterable[torch.Tensor]: for global_layer_id, local_layer_id in self.layer_offsets.items(): if self._is_local_mamba_layer(local_layer_id): continue # A layer group is a lifecycle, not a physical memory pool. # Differently sized attention buffers can share one lifecycle, # so scan every attention layer in this diagnostic path. - check_buffer(KVCacheManagerV2.get_buffers(self, global_layer_id)) + yield KVCacheManagerV2.get_buffers(self, global_layer_id) - for state_buffer in self.all_ssm_states: - check_buffer(state_buffer) - for state_buffer in self.all_conv_states: - check_buffer(state_buffer) - - torch.cuda.synchronize() - if some_checks_unavailable: - logger.warning( - "`torch.isnan` or `torch.isinf` is not implemented for a " - "configured cache dtype; related checks were skipped") - return bool(has_invalid_values) + yield from self.all_ssm_states + yield from self.all_conv_states def add_dummy_requests( self, @@ -3453,21 +3025,14 @@ def add_dummy_requests( num_extra_decoding_steps=num_extra_decoding_steps, draft_kv_cache_manager=draft_kv_cache_manager, ) - if requests: - self.requests.extend(requests) - if prepare_resource: - # self.requests may still contain transfer-pending requests - # from an older batch. Only the new dummy requests participate - # in this forward pass. - self._setup_state_indices(requests) + if requests and prepare_resource: + self._setup_state_indices(requests) return requests def free_resources(self, request: LlmRequest, pin_on_release: bool = False): kv_cache = self.kv_cache_map.get(request.py_request_id) if kv_cache is not None and kv_cache.is_active: self.try_commit_blocks(request, kv_cache) - if request in self.requests: - self.requests.remove(request) self._request_id_to_state_index.pop(request.py_request_id, None) super().free_resources(request, pin_on_release) @@ -3475,46 +3040,15 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): super().prepare_resources(scheduled_batch) if self.local_num_mamba_layers == 0: return - self.requests = (scheduled_batch.context_requests + - scheduled_batch.generation_requests) - self._setup_state_indices() + requests = (scheduled_batch.context_requests + + scheduled_batch.generation_requests) + self._setup_state_indices(requests) num_contexts = len(scheduled_batch.context_requests) - if num_contexts == 0: - return - ctx_slots = self.cuda_state_indices[:num_contexts].long() - if self._use_replay_state_update and self.prev_num_accepted_tokens is not None: - self.prev_num_accepted_tokens[ctx_slots] = 0 - self.cache_buf_idx[ctx_slots] = 0 - if self.old_x is not None: - self.old_x[:, ctx_slots] = 0 - if self.old_B is not None: - self.old_B[:, ctx_slots] = 0 - if self.old_dt is not None: - self.old_dt[:, ctx_slots] = 0 - if self.old_dA_cumsum is not None: - self.old_dA_cumsum[:, ctx_slots] = 0 - if self.mamba_ssm_rand_seed is not None: - self._seed_request_counter += 1 - counter = self._seed_request_counter - rank_offset = self._seed_rank_offset - host_slots = ctx_slots.cpu().tolist() - new_seeds = [ - _compute_deterministic_mamba_seed(counter, slot, rank_offset) - for slot in host_slots - ] - seed_tensor = torch.tensor( - new_seeds, - dtype=torch.int64, - device=self.mamba_ssm_rand_seed.device, - ) - self.mamba_ssm_rand_seed[ctx_slots] = seed_tensor + self._reset_context_mamba_slots(num_contexts) - def _setup_state_indices(self, - requests: Optional[List[LlmRequest]] = None - ) -> None: + def _setup_state_indices(self, requests: List[LlmRequest]) -> None: if self.local_num_mamba_layers == 0: return - requests = self.requests if requests is None else requests n = len(requests) assert n <= self._host_state_indices.shape[0], ( f"State-index batch size {n} exceeds max_batch_size " @@ -3557,9 +3091,6 @@ def get_state_indices(self, def get_max_resource_count(self) -> int: return self.max_batch_size - def is_speculative(self) -> bool: - return self.spec_config is not None - def update_mamba_states(self, attn_metadata: "AttentionMetadata", num_accepted_tokens: torch.Tensor, @@ -3579,23 +3110,13 @@ def update_mamba_states(self, src_state_indices = self.intermediate_state_indices[:num_gens] if self._use_replay_state_update: - # Mirror Mamba2Metadata's replay checkpoint predicate: only a - # checkpoint write resets PNAT and flips the active buffer. - slots = state_indices_d.long() - accepted = num_accepted_tokens[num_contexts:num_contexts + - num_gens].to( - self.prev_num_accepted_tokens. - dtype) - prev_num_accepted_tokens = self.prev_num_accepted_tokens[slots] - wrote_checkpoint = (prev_num_accepted_tokens + - self.replay_step_width - > self.replay_history_size) - self.prev_num_accepted_tokens[slots] = torch.where( - wrote_checkpoint, accepted, prev_num_accepted_tokens + accepted) - cache_buf_idx = self.cache_buf_idx[slots] - self.cache_buf_idx[slots] = torch.where(wrote_checkpoint, - 1 - cache_buf_idx, - cache_buf_idx) + replay_metadata = self.get_replay_state_update_metadata() + assert replay_metadata is not None + _advance_replay_state( + replay_metadata, + state_indices_d, + num_accepted_tokens[num_contexts:num_contexts + num_gens], + ) else: for layer_offset, dst in enumerate(self.all_ssm_states): _promote_mamba_state_triton( @@ -3615,82 +3136,6 @@ def update_mamba_states(self, state_indices_d, ) - def get_ssm_states(self, layer_idx: int) -> torch.Tensor: - return self.all_ssm_states[self.mamba_layer_offsets[layer_idx]] - - def get_conv_states(self, layer_idx: int) -> torch.Tensor: - return self.all_conv_states[self.mamba_layer_offsets[layer_idx]] - - def get_intermediate_ssm_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - if self.intermediate_ssm_states is None: - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.intermediate_ssm_states[layer_offset] - - def get_intermediate_conv_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - if self.intermediate_conv_states is None: - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.intermediate_conv_states[layer_offset] - - def mamba_layer_cache( - self, layer_idx: int - ) -> Union[PythonMambaCacheManager.State, - PythonMambaCacheManager.SpeculativeState, None]: - conv = self.get_conv_states(layer_idx) - ssm = self.get_ssm_states(layer_idx) - if self.spec_config is not None: - layer_offset = self.mamba_layer_offsets[layer_idx] - spec_kwargs = {} - if self.mamba_ssm_rand_seed is not None: - spec_kwargs['mamba_ssm_rand_seed'] = self.mamba_ssm_rand_seed - if self._use_replay_state_update: - spec_kwargs['old_x'] = self.old_x[layer_offset] - spec_kwargs['old_B'] = self.old_B[layer_offset] - spec_kwargs['old_dt'] = self.old_dt[layer_offset] - spec_kwargs['old_dA_cumsum'] = self.old_dA_cumsum[layer_offset] - spec_kwargs['cache_buf_idx'] = self.cache_buf_idx - spec_kwargs['prev_num_accepted_tokens'] = ( - self.prev_num_accepted_tokens) - else: - spec_kwargs['intermediate_ssm'] = self.intermediate_ssm_states[ - layer_offset] - return PythonMambaCacheManager.SpeculativeState( - conv=conv, - temporal=ssm, - intermediate_conv_window=self. - intermediate_conv_states[layer_offset], - **spec_kwargs, - ) - return PythonMambaCacheManager.State(conv=conv, temporal=ssm) - - def prepare_expect_snapshot_points(self, - requests: List[LlmRequest]) -> None: - """Set reusable Mamba snapshot boundaries before scheduling.""" - if not self.kv_cache_config.enable_block_reuse: - for request in requests: - request.expect_snapshot_points = [] - return - - state_config = self.kv_cache_config.mamba_state_config - interval = state_config.periodic_snapshot_interval - - for request in requests: - snapshot_points = set() - if interval is not None and interval > 0: - snapshot_points.update( - range(interval, request.prompt_len + 1, interval)) - for offset in state_config.additional_snapshot_offsets_from_start: - if offset <= request.prompt_len: - snapshot_points.add(offset) - for offset in state_config.additional_snapshot_offsets_from_end: - point = request.prompt_len - offset - if point > 0: - snapshot_points.add(point) - request.expect_snapshot_points = sorted(snapshot_points) - def _mark_context_position_as_history(self, request: LlmRequest, kv_cache) -> None: """Advance history without making later recurrent state reusable.""" @@ -3718,7 +3163,7 @@ def try_commit_blocks(self, request: LlmRequest, kv_cache=None) -> None: commit_limit = (min(max(snapshot_points), request.prompt_len) if snapshot_points else request.prompt_len) commit_end = min(request.context_current_position, commit_limit) - if (self._is_block_reuse_commit_boundary(request) + if (request.context_current_position in request.expect_snapshot_points and commit_end > kv_cache.num_committed_tokens): tokens = self._augment_tokens_for_block_reuse( request.get_tokens(DEFAULT_BEAM_INDEX), @@ -3765,36 +3210,6 @@ def update_context_resources(self, if request.context_remaining_length == 0: kv_cache.enable_swa_scratch_reuse = False - def _is_block_reuse_commit_boundary(self, request: LlmRequest) -> bool: - """Snapshot recurrent state only at an explicitly prepared point.""" - return request.context_current_position in request.expect_snapshot_points - - @property - def use_replay_state_update(self) -> bool: - return self.get_replay_state_update_metadata() is not None - - def get_replay_state_update_metadata( - self) -> Optional[ReplayStateUpdateMetadata]: - prev_num_accepted_tokens = getattr(self, 'prev_num_accepted_tokens', - None) - cache_buf_idx = getattr(self, 'cache_buf_idx', None) - if (not self._use_replay_state_update - or prev_num_accepted_tokens is None or cache_buf_idx is None - or self.replay_step_width is None - or self.replay_history_size is None): - return None - return ReplayStateUpdateMetadata( - prev_num_accepted_tokens=prev_num_accepted_tokens, - cache_buf_idx=cache_buf_idx, - replay_step_width=self.replay_step_width, - replay_history_size=self.replay_history_size) - - def get_mamba_ssm_cache_dtype(self) -> torch.dtype: - return self.ssm_state_dtype - - def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: - return getattr(self, 'mamba_ssm_rand_seed', None) - def shutdown(self): self.all_ssm_states = [] self.all_conv_states = [] diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 3e849cc8df12..9564226942be 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3836,6 +3836,18 @@ def validate_disk_cache_config(self): ) return self + @model_validator(mode='after') + def validate_mamba_snapshot_offsets(self) -> 'KvCacheConfig': + state_config = self.mamba_state_config + has_additional_snapshots = bool( + state_config.additional_snapshot_offsets_from_start + or state_config.additional_snapshot_offsets_from_end) + if (has_additional_snapshots and self.use_kv_cache_manager_v2 is False): + raise ValueError( + "kv_cache_config.mamba_state_config additional snapshot " + "offsets require kv_cache_config.use_kv_cache_manager_v2=True.") + return self + @field_validator('max_attention_window') @classmethod def validate_max_attention_window(cls, v: Optional[List[int]]): diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 7595f5b50dc7..076374b2721b 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -28,7 +28,9 @@ MambaRole, MixedMambaHybridCacheManager, PythonMambaCacheManager, + ReplayStateUpdateMetadata, V2MambaHybridCacheManager, + _advance_replay_state, _get_local_mamba_cache_layout, _get_mamba_hybrid_pool_size, _get_num_cuda_graph_padding_dummy_slots, @@ -43,7 +45,13 @@ from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._utils import torch_dtype_to_binding from tensorrt_llm.bindings.internal.batch_manager import LinearCacheType -from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MambaStateConfig, MTPDecodingConfig +from tensorrt_llm.llmapi.llm_args import ( + KvCacheConfig, + MambaStateConfig, + MTPDecodingConfig, + TorchLlmArgs, +) +from tensorrt_llm.llmapi.llm_utils import _resolve_kv_cache_manager_v2_auto from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import GpuCacheTierConfig, LayerId from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManager as RuntimeKVCacheManager @@ -51,6 +59,27 @@ skip_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_advance_replay_state_uses_checkpoint_predicate_and_skips_dummies(): + metadata = ReplayStateUpdateMetadata( + prev_num_accepted_tokens=torch.tensor([11, 12, 13], dtype=torch.int32), + cache_buf_idx=torch.tensor([0, 1, 1], dtype=torch.int32), + replay_step_width=5, + replay_history_size=16, + ) + + _advance_replay_state( + metadata, + state_indices=torch.tensor([0, 1, 2], dtype=torch.int32), + accepted_tokens=torch.tensor([2, 3, 4], dtype=torch.int32), + is_dummy_request=torch.tensor([False, False, True]), + ) + + # Equality does not write a checkpoint; overflow does. Dummy slots do not + # advance either piece of replay bookkeeping. + assert metadata.prev_num_accepted_tokens.tolist() == [13, 3, 13] + assert metadata.cache_buf_idx.tolist() == [0, 0, 1] + + def test_cuda_graph_padding_dummy_slot_count_tracks_reachable_draft_lengths(): assert _get_num_cuda_graph_padding_dummy_slots(None, 64) == 1 @@ -243,20 +272,37 @@ def test_hybrid_cache_manager_factory_v2_preference_does_not_select_v2( ) +@pytest.mark.parametrize( + ("field", "offsets"), + [ + ("additional_snapshot_offsets_from_start", [128]), + ("additional_snapshot_offsets_from_end", [0]), + ], +) def test_hybrid_cache_manager_factory_requires_v2_for_explicit_snapshots( monkeypatch, + field, + offsets, ): monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + kv_cache_config = KvCacheConfig( + enable_block_reuse=False, + mamba_state_config=MambaStateConfig(**{field: offsets}), + use_kv_cache_manager_v2="auto", + ) + llm_args = TorchLlmArgs( + model="/tmp/dummy_model", + kv_cache_config=kv_cache_config, + ) + + assert _resolve_kv_cache_manager_v2_auto(llm_args, {}) is False + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False with pytest.raises(ValueError, match="use_kv_cache_manager_v2=True"): get_kv_cache_manager_cls( _hybrid_model_config(), - KvCacheConfig( - enable_block_reuse=False, - mamba_state_config=MambaStateConfig(additional_snapshot_offsets_from_end=[0]), - use_kv_cache_manager_v2=False, - ), + llm_args.kv_cache_config, ) @@ -275,8 +321,17 @@ def test_hybrid_cache_manager_factory_requires_snapshot_policy(monkeypatch): ) +@pytest.mark.parametrize( + ("field", "offsets"), + [ + ("additional_snapshot_offsets_from_start", [128]), + ("additional_snapshot_offsets_from_end", [0]), + ], +) def test_hybrid_cache_manager_factory_routes_explicit_snapshots_to_v2( monkeypatch, + field, + offsets, ): monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) @@ -288,7 +343,7 @@ def test_hybrid_cache_manager_factory_routes_explicit_snapshots_to_v2( enable_block_reuse=True, mamba_state_config=MambaStateConfig( periodic_snapshot_interval=0, - additional_snapshot_offsets_from_end=[0], + **{field: offsets}, ), use_kv_cache_manager_v2=True, ), @@ -1232,8 +1287,10 @@ def test_cpp_hybrid_prepare_expect_snapshot_points_clears_when_reuse_disabled(): @pytest.mark.parametrize("interval", [0, -64, None]) def test_cpp_hybrid_prepare_expect_snapshot_points_clears_for_invalid_interval(interval): mgr = object.__new__(CppMambaHybridCacheManager) - mgr.kv_cache_config = SimpleNamespace(enable_block_reuse=True) - mgr.linear_attention_metadata = SimpleNamespace(states_snapshot_interval=interval) + mgr.kv_cache_config = SimpleNamespace( + enable_block_reuse=True, + mamba_state_config=SimpleNamespace(periodic_snapshot_interval=interval), + ) request = SimpleNamespace(prompt_len=150, expect_snapshot_points=[64]) mgr.prepare_expect_snapshot_points([request]) @@ -1387,6 +1444,7 @@ def _build_v2_hybrid_with_mamba_layer( enable_partial_reuse=True, enable_attention_dp=False, enable_swa_scratch_reuse=False, + dtype=DataType.HALF, ): """Construct a real V2MambaHybridCacheManager.""" mamba_mask = [True] * num_mamba_layers + [False] * num_attention_layers @@ -1404,6 +1462,7 @@ def _build_v2_hybrid_with_mamba_layer( enable_block_reuse=enable_block_reuse, enable_partial_reuse=enable_partial_reuse, enable_swa_scratch_reuse=enable_swa_scratch_reuse, + dtype="nvfp4" if dtype == DataType.NVFP4 else "auto", ) return V2MambaHybridCacheManager( mamba_d_state=8, @@ -1428,6 +1487,7 @@ def _build_v2_hybrid_with_mamba_layer( layer_mask=attn_mask, vocab_size=1024, use_replay_state_update=use_replay_state_update, + dtype=dtype, ) @@ -1496,6 +1556,19 @@ def test_v2_hybrid_allocates_mamba_state_and_dummy_indices(): mgr.shutdown() +@skip_no_cuda +def test_v2_hybrid_nvfp4_page_table_omits_ssm_block_scales(): + mgr = _build_v2_hybrid_with_mamba_layer(dtype=DataType.NVFP4) + try: + ssm_pool_id = mgr.impl.get_layer_group_id(LayerId(0)) + attention_pool_id = mgr.impl.get_layer_group_id(LayerId(1)) + + assert torch.count_nonzero(mgr.kv_cache_pool_pointers[ssm_pool_id, :, 1]) == 0 + assert torch.count_nonzero(mgr.kv_cache_pool_pointers[attention_pool_id, :, 1]) > 0 + finally: + mgr.shutdown() + + @skip_no_cuda def test_v2_hybrid_supports_pure_mamba_pp_rank(): mgr = _build_v2_hybrid_with_mamba_layer( @@ -1555,9 +1628,15 @@ def test_v2_hybrid_invalid_check_scans_distinct_attention_pools(): second_attention_buffer = mgr.get_buffers(2) assert first_attention_buffer.data_ptr() != second_attention_buffer.data_ptr() - second_attention_buffer.flatten()[0] = torch.nan - - assert mgr.check_invalid_values_in_kv_cache() + for buffer in ( + second_attention_buffer, + mgr.all_ssm_states[0], + mgr.all_conv_states[0], + ): + buffer.flatten()[0] = torch.nan + assert mgr.check_invalid_values_in_kv_cache() + assert mgr.check_invalid_values_in_kv_cache(fill_with_zero=True) + assert not torch.isnan(buffer).any() finally: mgr.shutdown() @@ -1567,18 +1646,19 @@ def test_v2_hybrid_invalid_check_scans_distinct_attention_pools(): def test_v2_hybrid_dummy_indices_keep_cuda_buffer_address(max_batch_size): mgr = _build_v2_hybrid_with_mamba_layer(max_batch_size=max_batch_size, enable_attention_dp=True) try: - stale_request = mgr.add_dummy_requests([123], token_nums=[8], is_gen=False)[0] + request_ids = list(range(100, 100 + max_batch_size)) + mgr.add_dummy_requests( + request_ids, + token_nums=[8] * max_batch_size, + is_gen=False, + ) state_indices_ptr = mgr.cuda_state_indices.data_ptr() - # Model a transfer-pending prior batch. The new ADP dummy is the only - # request participating in the next forward pass. - mgr.requests = [stale_request] * max_batch_size new_requests = mgr.add_dummy_requests( [ATTENTION_DP_DUMMY_REQUEST_ID], token_nums=[8], is_gen=False ) assert len(new_requests) == 1 - assert len(mgr.requests) == max_batch_size + 1 expected_capacity = max_batch_size + mgr._num_reserved_dummy_slots assert mgr.cuda_state_indices.shape[0] == expected_capacity assert mgr._host_state_indices.shape[0] == expected_capacity @@ -1648,7 +1728,9 @@ def test_v2_hybrid_free_resources_drops_stale_state_index_mapping(): request_id = request.py_request_id assert request_id in mgr._request_id_to_state_index - mgr.requests.clear() + # Move state-index preparation to another request before freeing the + # older one, as happens when an asynchronous transfer finishes late. + mgr.add_dummy_requests([456], token_nums=[8], is_gen=False) mgr.free_resources(request) assert request_id not in mgr._request_id_to_state_index @@ -1775,6 +1857,24 @@ def test_cpp_hybrid_replay_buffers_size_by_tokens_per_gen_step(): mgr.shutdown() +@skip_no_cuda +def test_v2_hybrid_intermediate_states_size_by_tokens_per_gen_step(): + spec_config = _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5) + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=4, + spec_config=spec_config, + ) + try: + assert mgr.intermediate_ssm_states.shape[2] == 5 + assert mgr.intermediate_conv_states.shape[2] == 5 + layer_cache = mgr.mamba_layer_cache(0) + assert layer_cache.intermediate_ssm.data_ptr() == ( + mgr.intermediate_ssm_states[0].data_ptr() + ) + finally: + mgr.shutdown() + + @skip_no_cuda def test_v2_hybrid_replay_buffers_size_by_tokens_per_gen_step(): spec_config = _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 807d54baeff9..9fcee7937857 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -675,6 +675,40 @@ def test_MambaStateConfig_rejects_invalid_snapshot_offsets(field, value): MambaStateConfig(**{field: value}) +@pytest.mark.parametrize( + ("field", "offsets"), + [ + ("additional_snapshot_offsets_from_start", [128]), + ("additional_snapshot_offsets_from_end", [0]), + ], +) +def test_KvCacheConfig_requires_v2_for_additional_snapshot_offsets( + field, offsets): + state_config = MambaStateConfig(**{field: offsets}) + + with pytest.raises(ValidationError, match="use_kv_cache_manager_v2=True"): + KvCacheConfig( + mamba_state_config=state_config, + use_kv_cache_manager_v2=False, + ) + + for use_v2 in (True, "auto"): + config = KvCacheConfig( + mamba_state_config=state_config, + use_kv_cache_manager_v2=use_v2, + ) + assert getattr(config.mamba_state_config, field) == offsets + + +def test_KvCacheConfig_allows_periodic_snapshots_with_v1(): + config = KvCacheConfig( + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=64), + use_kv_cache_manager_v2=False, + ) + + assert config.mamba_state_config.periodic_snapshot_interval == 64 + + def test_KvCacheConfig_rejects_removed_mamba_interval_field(): with pytest.raises(ValidationError, match="extra_forbidden"): KvCacheConfig(mamba_state_cache_interval=64) From 3aec7fd2b8c55dd348cf059d6cacfe776361f1b0 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 21 Jul 2026 18:08:52 +0800 Subject: [PATCH 03/22] [TRTLLM-11875][fix] Disable periodic Mamba snapshots by default (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- docs/source/features/kvcache.md | 12 ++++++++---- tensorrt_llm/llmapi/llm_args.py | 5 +++-- .../_torch/executor/test_mamba_cache_manager.py | 14 +++++++++++--- tests/unittest/llmapi/test_llm_args.py | 2 +- 4 files changed, 23 insertions(+), 10 deletions(-) diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index f459143ecfd7..4768dfd9bc15 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -79,10 +79,14 @@ scheduler_config: Hybrid Mamba models must retain the recurrent Mamba state together with the attention KV prefix. Snapshot policy is grouped under `kv_cache_config.mamba_state_config`. `periodic_snapshot_interval` controls -periodic boundaries; set it to `0` to disable them. The interval is accepted -through this nested configuration in the Python API. YAML/JSON configuration -files also accept the deprecated `kv_cache_config.mamba_state_cache_interval` -key and migrate it to the nested field while loading. The prototype +periodic boundaries. They are disabled by default; set the interval to a +positive value to enable them. The interval is accepted through this nested +configuration in the Python API. YAML/JSON configuration files also accept the +deprecated `kv_cache_config.mamba_state_cache_interval` key and migrate it to +the nested field while loading. Enabling block reuse for a hybrid Mamba model +requires either a positive periodic interval or an explicit V2 snapshot offset; +otherwise, configuration validation rejects the unsupported reuse policy. The +prototype `additional_snapshot_offsets_from_start` and `additional_snapshot_offsets_from_end` options add fixed boundaries. Start offsets count tokens from the beginning of the prompt. End offsets count diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 9564226942be..a8f7ce63a598 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3466,10 +3466,11 @@ class MambaStateConfig(StrictBaseModel): """Configuration for reusable Mamba recurrent-state snapshots.""" periodic_snapshot_interval: NonNegativeInt = Field( - default=256, + default=0, description= "The number of tokens between periodic snapshots in the Mamba " - "prefix cache. Set to 0 to disable periodic snapshots.") + "prefix cache. Periodic snapshots are disabled by default; set this " + "to a positive value to enable them.") additional_snapshot_offsets_from_start: List[_StrictPositiveInt] = Field( default_factory=list, diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 076374b2721b..e28046cd8bfb 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -157,6 +157,7 @@ def test_hybrid_cache_manager_factory_honors_v2_setting( monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) kv_cache_config = KvCacheConfig( enable_block_reuse=enable_block_reuse, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), use_kv_cache_manager_v2=use_v2, ) @@ -252,6 +253,7 @@ def test_hybrid_cache_manager_factory_rejects_cpp_preference_with_explicit_v2( _hybrid_model_config(), KvCacheConfig( enable_block_reuse=True, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), use_kv_cache_manager_v2=True, ), ) @@ -266,7 +268,10 @@ def test_hybrid_cache_manager_factory_v2_preference_does_not_select_v2( assert ( get_kv_cache_manager_cls( _hybrid_model_config(), - KvCacheConfig(use_kv_cache_manager_v2=False), + KvCacheConfig( + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + use_kv_cache_manager_v2=False, + ), ) is CppMambaHybridCacheManager ) @@ -377,7 +382,10 @@ def test_hybrid_cache_manager_factory_rejects_mixed_override_with_reuse( with pytest.raises(ValueError, match=expected_error): get_kv_cache_manager_cls( _hybrid_model_config(), - KvCacheConfig(enable_block_reuse=True), + KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + ), ) @@ -1036,7 +1044,7 @@ def test_v2_hybrid_estimator_accounts_for_ssm_slot_attention_bound( max_batch_size=4, kv_cache_config=KvCacheConfig(), spec_config=spec_config, - ) == (12, expected_intercept) + ) == (11, expected_intercept) def test_v2_hybrid_estimator_reserves_attention_for_each_lineage(monkeypatch): diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 9fcee7937857..64a8a6b94aff 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -641,7 +641,7 @@ def test_MambaStateConfig_defaults_use_independent_lists(): first = MambaStateConfig() second = MambaStateConfig() - assert first.periodic_snapshot_interval == 256 + assert first.periodic_snapshot_interval == 0 first.additional_snapshot_offsets_from_start.append(128) first.additional_snapshot_offsets_from_end.append(0) From 6558b7835602b6ae773127b5af703846680d759d Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:27:18 +0800 Subject: [PATCH 04/22] [None][feat] Support V2 Mamba disaggregated serving (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../disaggregation/native/mixers/ssm/peer.py | 60 +++- .../disaggregation/resource/kv_extractor.py | 102 +++++- .../_torch/disaggregation/resource/page.py | 56 +++ .../_torch/disaggregation/resource/utils.py | 38 ++- .../_torch/disaggregation/transceiver.py | 14 +- .../_torch/models/modeling_nemotron_h.py | 9 +- .../_torch/models/modeling_qwen3_5.py | 9 +- .../_torch/models/modeling_qwen3_next.py | 9 +- tensorrt_llm/_torch/pyexecutor/_util.py | 71 +++- .../_torch/pyexecutor/kv_cache_transceiver.py | 34 +- .../_torch/pyexecutor/mamba_cache_manager.py | 35 +- .../_torch/pyexecutor/py_executor_creator.py | 21 +- .../test_disagg_inflight_cancel_gate.py | 57 ++++ .../executor/test_mamba_cache_manager.py | 201 +++++++++-- .../unittest/disaggregated/test_extractor.py | 126 ++++++- .../disaggregated/test_mamba_transfer.py | 318 ++++++++++++++++-- 16 files changed, 1035 insertions(+), 125 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py index e8c8816ba4c8..81feb03a25a7 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py @@ -1,3 +1,17 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from typing import Dict, List, Optional, Tuple import numpy as np @@ -337,14 +351,28 @@ def _build_layer_ptrs( layer_offsets: Dict[int, int], overlapping_layers: List[int], slot: int, + layer_slot0_addresses: Optional[Dict[int, int]] = None, ) -> np.ndarray: - """Build per-layer pointers for a given pool (conv or ssm) and slot.""" + """Build per-layer pointers for a given pool (conv or SSM) and slot. + + V1 stores states layer-major, so its layer base is derived from the + Mamba-local layer offset. V2 stores buffers inside coalesced slot-major + pools; its manager-provided slot-0 addresses preserve the per-layer + offsets within that physical slot. + """ ptrs = [] + slot_stride_bytes = pool.slot_stride_bytes + assert slot_stride_bytes is not None for glid in overlapping_layers: - lid = layer_offsets[glid] - ptrs.append( - pool.base_address + lid * pool.num_slots * pool.slot_bytes + slot * pool.slot_bytes - ) + if layer_slot0_addresses is not None: + ptrs.append(layer_slot0_addresses[glid] + slot * slot_stride_bytes) + else: + lid = layer_offsets[glid] + ptrs.append( + pool.base_address + + lid * pool.num_slots * pool.slot_bytes + + slot * pool.slot_bytes + ) return np.array(ptrs, dtype=np.int64) @staticmethod @@ -430,11 +458,29 @@ def build_mamba_frags( (self_mlg.conv_states, peer_mlg.conv_states, True), (self_mlg.ssm_states, peer_mlg.ssm_states, False), ]: + self_layer_slot0_addresses = ( + self_mlg.conv_layer_slot0_addresses + if is_conv + else self_mlg.ssm_layer_slot0_addresses + ) + peer_layer_slot0_addresses = ( + peer_mlg.conv_layer_slot0_addresses + if is_conv + else peer_mlg.ssm_layer_slot0_addresses + ) src_ptrs = MambaPolicy._build_layer_ptrs( - self_pool, self_mlg.mamba_layer_offsets, overlapping_layers, src_slot + self_pool, + self_mlg.mamba_layer_offsets, + overlapping_layers, + src_slot, + self_layer_slot0_addresses, ) dst_ptrs = MambaPolicy._build_layer_ptrs( - peer_pool, peer_mlg.mamba_layer_offsets, overlapping_layers, dst_slot + peer_pool, + peer_mlg.mamba_layer_offsets, + overlapping_layers, + dst_slot, + peer_layer_slot0_addresses, ) src_region = SpecRegion( diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index 92f0a9c2f4ad..ddf876bd7dbd 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -1,3 +1,17 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from typing import Dict, List import numpy as np @@ -21,7 +35,10 @@ PoolView, ) from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( + MambaHybridCacheManager, + V2MambaHybridCacheManager, +) from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import get_size_in_bytes, nvtx_range from tensorrt_llm.bindings import DataType @@ -73,10 +90,12 @@ def extract( base_ptr = pool.base_address block_size = pool.slot_bytes + block_stride = pool.slot_stride_bytes + assert block_stride is not None # KV cache: filter out invalid block_ids (BAD_PAGE_INDEX = -1) valid = region_ids >= 0 - ptrs = base_ptr + block_size * region_ids[valid] + ptrs = base_ptr + block_stride * region_ids[valid] memory = MemRegionGroup(ptrs=ptrs, bytes_per_region=block_size) return SpecRegion(memory=memory) @@ -110,7 +129,8 @@ def _build_layer_group_for_mamba( ) # Per-section bytes for conv_state and per-head bytes for ssm_state. - # conv_state layout: [x: d_inner/tp | B: ng*ds/tp | C: ng*ds/tp] x (d_conv-1) + # The section ordering is supplied by the cache manager because Mamba2 + # uses [x | B | C], while GDN uses [Q | K | V]. # ssm_state layout: (nheads/tp, head_dim, d_state) d_conv_m1 = conv_state.shape[3] conv_elem_size = conv_state.element_size() @@ -132,6 +152,70 @@ def _build_layer_group_for_mamba( ) +def _slot_stride_bytes(tensor) -> int: + return int(tensor.stride(0) * tensor.element_size()) + + +def _build_layer_group_for_v2_mamba( + manager: V2MambaHybridCacheManager, pool_group_idx: int +) -> MambaLayerGroup: + mamba_layer_offsets = { + int(global_layer_id): int(local_layer_id) + for global_layer_id, local_layer_id in manager.mamba_layer_offsets.items() + } + + first_conv_state = manager.all_conv_states[0] + first_ssm_state = manager.all_ssm_states[0] + conv_slot_stride_bytes = _slot_stride_bytes(first_conv_state) + ssm_slot_stride_bytes = _slot_stride_bytes(first_ssm_state) + conv_slot_bytes = int(first_conv_state[0].numel() * first_conv_state.element_size()) + ssm_slot_bytes = int(first_ssm_state[0].numel() * first_ssm_state.element_size()) + num_slots = int(first_ssm_state.shape[0]) + + # V2 coalesces equal-size buffers into slot-major physical pools. The + # SHARED tensor bases include each layer/role's offset within slot 0, while + # stride(0) is the distance to the same buffer in the next physical slot. + # Preserve both pieces: V1's layer-major ``layer * num_slots`` formula does + # not describe this layout. + conv_layer_slot0_addresses = { + int(global_layer_id): int(manager.all_conv_states[offset].data_ptr()) + for global_layer_id, offset in mamba_layer_offsets.items() + } + ssm_layer_slot0_addresses = { + int(global_layer_id): int(manager.all_ssm_states[offset].data_ptr()) + for global_layer_id, offset in mamba_layer_offsets.items() + } + + d_conv_m1 = manager.conv_state_shape[1] + conv_elem_size = first_conv_state.element_size() + _, head_dim, d_state = manager.ssm_state_shape + conv_section_bytes = [dim * d_conv_m1 * conv_elem_size for dim in manager.conv_section_dims] + + ssm_elem_size = first_ssm_state.element_size() + ssm_bytes_per_head = head_dim * d_state * ssm_elem_size + + return MambaLayerGroup( + pool_group_idx=pool_group_idx, + mamba_layer_offsets=mamba_layer_offsets, + conv_states=PhysicalPool( + base_address=int(first_conv_state.data_ptr()), + slot_bytes=conv_slot_bytes, + num_slots=num_slots, + slot_stride_bytes=conv_slot_stride_bytes, + ), + ssm_states=PhysicalPool( + base_address=int(first_ssm_state.data_ptr()), + slot_bytes=ssm_slot_bytes, + num_slots=num_slots, + slot_stride_bytes=ssm_slot_stride_bytes, + ), + conv_section_bytes=conv_section_bytes, + ssm_bytes_per_head=ssm_bytes_per_head, + conv_layer_slot0_addresses=conv_layer_slot0_addresses, + ssm_layer_slot0_addresses=ssm_layer_slot0_addresses, + ) + + def build_page_table(kv_cache_manager: KVCacheManager) -> KVCachePageTable: """Build a KVCachePageTable from a KVCacheManager (V1).""" if kv_cache_manager.dtype == DataType.NVFP4: @@ -339,6 +423,14 @@ def _window_size_for_layer(internal_layer_id: int): for variant in pg_desc.slot_desc.variants: layer_group_id = int(variant.layer_group_id) all_internal_layer_ids = list(manager.impl.layer_grouping[layer_group_id]) + if isinstance(manager, V2MambaHybridCacheManager) and any( + manager._is_local_mamba_layer(int(layer_id)) for layer_id in all_internal_layer_ids + ): + layer_groups_by_id[layer_group_id] = _build_layer_group_for_v2_mamba( + manager, storage_pg_to_list_idx[storage_pg_idx] + ) + continue + all_global_layer_ids = _compute_global_layer_ids(manager, layer_group_id) local_layers = [ @@ -392,7 +484,9 @@ def _window_size_for_layer(internal_layer_id: int): raise ValueError(f"Missing V2 layer group descriptor for layer group {layer_group_id}") layer_groups.append(layer_group) - if isinstance(manager, MambaHybridCacheManager): + if isinstance(manager, MambaHybridCacheManager) and not isinstance( + manager, V2MambaHybridCacheManager + ): mamba_layer_group_idx = len(pool_groups) mamba_layer_group = _build_layer_group_for_mamba(manager, mamba_layer_group_idx) layer_groups.append(mamba_layer_group) diff --git a/tensorrt_llm/_torch/disaggregation/resource/page.py b/tensorrt_llm/_torch/disaggregation/resource/page.py index 81514c06d6a0..28537ff9a5bb 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/page.py +++ b/tensorrt_llm/_torch/disaggregation/resource/page.py @@ -1,3 +1,17 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from __future__ import annotations from dataclasses import dataclass, field @@ -48,12 +62,24 @@ class PhysicalPool: base_address: int # uint64 slot_bytes: int num_slots: int + # Distance between the starts of adjacent slots. Most pools are densely + # packed, so the stride defaults to the transferable payload size. V2 + # Mamba views point into a coalesced physical slot whose stride can be + # larger than the state payload described by ``slot_bytes``. + slot_stride_bytes: Optional[int] = None + + def __post_init__(self) -> None: + if self.slot_stride_bytes is None: + self.slot_stride_bytes = self.slot_bytes + if self.slot_stride_bytes < self.slot_bytes: + raise ValueError("slot_stride_bytes must be greater than or equal to slot_bytes") def to_dict(self) -> dict: return { "base_address": int(self.base_address), "slot_bytes": int(self.slot_bytes), "num_slots": int(self.num_slots), + "slot_stride_bytes": int(self.slot_stride_bytes), } @staticmethod @@ -62,6 +88,11 @@ def from_dict(data: dict) -> "PhysicalPool": base_address=int(data["base_address"]), slot_bytes=int(data["slot_bytes"]), num_slots=int(data["num_slots"]), + slot_stride_bytes=( + int(data["slot_stride_bytes"]) + if data.get("slot_stride_bytes") is not None + else None + ), ) @@ -204,6 +235,11 @@ class MambaLayerGroup(LayerGroup): ssm_states: Optional[PhysicalPool] = None conv_section_bytes: Optional[List[int]] = None ssm_bytes_per_head: Optional[int] = None + # V2 pools are slot-major and may coalesce several layer/role buffers into + # one physical slot. These are the manager-provided buffer offsets within + # slot 0; they cannot be derived with V1's layer-major pointer formula. + conv_layer_slot0_addresses: Optional[Dict[int, int]] = None + ssm_layer_slot0_addresses: Optional[Dict[int, int]] = None def to_dict(self) -> dict: return { @@ -213,12 +249,24 @@ def to_dict(self) -> dict: "ssm_states": self.ssm_states.to_dict(), "conv_section_bytes": self.conv_section_bytes, "ssm_bytes_per_head": self.ssm_bytes_per_head, + "conv_layer_slot0_addresses": { + int(k): int(v) for k, v in (self.conv_layer_slot0_addresses or {}).items() + } + if self.conv_layer_slot0_addresses is not None + else None, + "ssm_layer_slot0_addresses": { + int(k): int(v) for k, v in (self.ssm_layer_slot0_addresses or {}).items() + } + if self.ssm_layer_slot0_addresses is not None + else None, } @classmethod def from_dict(cls, data: dict) -> "MambaLayerGroup": conv_section_bytes = data.get("conv_section_bytes") ssm_bytes_per_head = data.get("ssm_bytes_per_head") + conv_layer_slot0_addresses = data.get("conv_layer_slot0_addresses") + ssm_layer_slot0_addresses = data.get("ssm_layer_slot0_addresses") return cls( pool_group_idx=int(data["pool_group_idx"]), mamba_layer_offsets={int(k): int(v) for k, v in data["mamba_layer_offsets"].items()}, @@ -228,6 +276,14 @@ def from_dict(cls, data: dict) -> "MambaLayerGroup": if conv_section_bytes is not None else None, ssm_bytes_per_head=int(ssm_bytes_per_head) if ssm_bytes_per_head is not None else None, + conv_layer_slot0_addresses={ + int(k): int(v) for k, v in conv_layer_slot0_addresses.items() + } + if conv_layer_slot0_addresses is not None + else None, + ssm_layer_slot0_addresses={int(k): int(v) for k, v in ssm_layer_slot0_addresses.items()} + if ssm_layer_slot0_addresses is not None + else None, ) diff --git a/tensorrt_llm/_torch/disaggregation/resource/utils.py b/tensorrt_llm/_torch/disaggregation/resource/utils.py index 21c4d98bd2aa..874911f81fd9 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/utils.py +++ b/tensorrt_llm/_torch/disaggregation/resource/utils.py @@ -1,3 +1,17 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from __future__ import annotations from typing import Dict, List, Set @@ -10,7 +24,7 @@ def get_pool_bytes(pool: PhysicalPool) -> int: - """Total bytes across all slots in this pool.""" + """Total transferable payload bytes across all slots in this pool.""" return pool.slot_bytes * pool.num_slots @@ -18,7 +32,8 @@ def get_slot_address(pool: PhysicalPool, slot_id: int) -> int: """Base address of *slot_id*.""" if slot_id >= pool.num_slots: raise ValueError(f"slot_id {slot_id} >= num_slots {pool.num_slots}") - return pool.base_address + slot_id * pool.slot_bytes + assert pool.slot_stride_bytes is not None + return pool.base_address + slot_id * pool.slot_stride_bytes # ------------------------------------------------------------------------- @@ -117,9 +132,22 @@ def get_unique_pool_memory_descs( pool_counter = 0 for lg_idx, lg in enumerate(page_table.layer_groups): if isinstance(lg, MambaLayerGroup): - num_mamba_layers = len(lg.mamba_layer_offsets) - for pool in [lg.conv_states, lg.ssm_states]: - pool_size = num_mamba_layers * pool.num_slots * pool.slot_bytes + is_v2_layout = ( + lg.conv_layer_slot0_addresses is not None + or lg.ssm_layer_slot0_addresses is not None + ) + if is_v2_layout: + pools_and_sizes = [ + (pool, get_pool_bytes(pool)) + for pool in page_table.pool_groups[int(lg.pool_group_idx)].pools + ] + else: + num_mamba_layers = len(lg.mamba_layer_offsets) + pools_and_sizes = [ + (pool, num_mamba_layers * pool.num_slots * pool.slot_bytes) + for pool in [lg.conv_states, lg.ssm_states] + ] + for pool, pool_size in pools_and_sizes: pool_key = (pool.base_address, pool_size) if pool_key not in unique_pools: unique_pools[pool_key] = pool_counter diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 4fe7fa28e50f..3bcc7ec72d6f 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -30,7 +30,10 @@ from tensorrt_llm._torch.distributed.communicator import Distributed from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import KvCacheTransceiver from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( + MambaHybridCacheManager, + V2MambaHybridCacheManager, +) from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import nvtx_range from tensorrt_llm.bindings import LlmRequestState @@ -137,7 +140,9 @@ def _init_sync_policy(self): def _exchange_rank_info(self): endpoints = cast(list, self._dist.allgather(self._transfer_worker.sender_endpoint)) layer_num = len(self._kv_cache_manager.pp_layers) - if isinstance(self._kv_cache_manager, MambaHybridCacheManager): + if isinstance(self._kv_cache_manager, MambaHybridCacheManager) and not isinstance( + self._kv_cache_manager, V2MambaHybridCacheManager + ): layer_num += len(self._kv_cache_manager._impl.mamba_layer_offsets) layer_num_per_pp = cast(list, getattr(self._dist, "pp_allgather")(layer_num)) self._transfer_worker.populate_instance_and_rank_info( @@ -230,7 +235,10 @@ def _create_kv_slice(self, req: LlmRequest) -> KVSlice: groups.append(block_ids) mamba_state_index = None - if isinstance(self._kv_cache_manager, MambaHybridCacheManager): + if isinstance(self._kv_cache_manager, V2MambaHybridCacheManager): + if self._kv_cache_manager.local_num_mamba_layers > 0: + mamba_state_index = self._kv_cache_manager.get_state_indices([req.py_request_id])[0] + elif isinstance(self._kv_cache_manager, MambaHybridCacheManager): mamba_state_index = self._kv_cache_manager.mamba_cache_index[req.py_request_id] return KVSlice( diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 5c2e60c4706a..2938c88507b1 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -17,7 +17,7 @@ import re from contextlib import contextmanager from dataclasses import replace -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal import torch @@ -991,6 +991,13 @@ def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: # is supported for Mamba/SSM-based models return {"kv_cache_config": {"enable_block_reuse": False}} + @classmethod + def get_preferred_transceiver_runtime(cls, + pretrained_config: object + | None = None) -> Literal["PYTHON"]: + """Use the Python transceiver for hybrid-state transfers.""" + return "PYTHON" + @staticmethod def lora_config(model_dir: str): """Nemotron-H-specific LoRA configuration. diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_5.py b/tensorrt_llm/_torch/models/modeling_qwen3_5.py index 6fc43129a353..4bd9de60c284 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_5.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_5.py @@ -15,7 +15,7 @@ import re from types import SimpleNamespace -from typing import Dict, List +from typing import Dict, List, Literal import torch from transformers import PretrainedConfig @@ -674,6 +674,13 @@ def get_model_defaults(cls, llm_args): # would silently fall back to the global default (block reuse on). return Qwen3NextForCausalLM.get_model_defaults(llm_args) + @classmethod + def get_preferred_transceiver_runtime( + cls, pretrained_config: object | None = None + ) -> Literal["PYTHON"]: + """Match the hybrid text decoder's Python disaggregated route.""" + return "PYTHON" + def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, **kwargs): kwargs["vision_model_class"] = Qwen3VisionModel kwargs["disable_fuse_rope"] = kwargs.get("disable_fuse_rope", False) diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index 0d91b4ebad74..87925dbebe06 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py @@ -18,7 +18,7 @@ import copy import os from types import SimpleNamespace -from typing import TYPE_CHECKING, Dict, List, Optional +from typing import TYPE_CHECKING, Dict, List, Literal, Optional import torch @@ -992,6 +992,13 @@ def get_model_defaults(cls, llm_args: 'TorchLlmArgs') -> dict: # is supported for Mamba/SSM-based models return {"kv_cache_config": {"enable_block_reuse": False}} + @classmethod + def get_preferred_transceiver_runtime(cls, + pretrained_config: object + | None = None) -> Literal["PYTHON"]: + """Use the Python transceiver for hybrid-state transfers.""" + return "PYTHON" + def load_weights(self, weights: dict, weight_mapper: BaseWeightMapper, diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 032e8e1881fa..23b59ff2f1f5 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -88,6 +88,22 @@ def _non_hybrid_kv_cache_manager_cls(config, kv_cache_config: KvCacheConfig): return KVCacheManagerV2 if needs_v2 else KVCacheManager +def _resolve_disagg_transceiver_route( + cache_transceiver_config: Optional[CacheTransceiverConfig], +) -> tuple[Optional[str], Optional[str]]: + """Return the effective backend and runtime used for manager routing.""" + if cache_transceiver_config is None: + return None, None + + backend, _ = cache_transceiver_config._resolve_default_backend() + runtime = cache_transceiver_config.transceiver_runtime + if runtime == "auto": + # Model loading normally resolves ``auto``. Paths that skip model + # defaults use the global C++ fallback, matching transceiver creation. + runtime = None + return backend, runtime + + def get_kv_cache_manager_cls( model_config: ModelConfig, kv_cache_config: KvCacheConfig, @@ -101,8 +117,14 @@ def get_kv_cache_manager_cls( unified-pool default. V1 is the default for hybrid Mamba models. V2 is selected only by an - explicit ``kv_cache_config.use_kv_cache_manager_v2=True`` and is rejected - for disaggregated serving until its transceiver/page-table adapter exists. + explicit ``kv_cache_config.use_kv_cache_manager_v2=True``. In + disaggregated serving, V2 additionally requires the Python transceiver + with the NIXL backend. Unsupported explicit V2 routes fail rather than + falling back to a different manager. + + Env-var overrides: + * ``TRTLLM_USE_PY_MAMBA=1`` — Mixed manager in aggregated serving. + * ``TLLM_MAMBA_MANAGER_PREFERENCE`` — explicit manager preference. """ config = model_config.pretrained_config sparse_attn_config = model_config.sparse_attention_config @@ -131,12 +153,6 @@ def get_kv_cache_manager_cls( or state_config.additional_snapshot_offsets_from_end) use_v2 = kv_cache_config.use_kv_cache_manager_v2 is True - if is_disagg and use_v2: - raise ValueError( - "KV cache manager V2 for hybrid Mamba models is not " - "supported with disaggregated serving. Set " - "use_kv_cache_manager_v2=False or 'auto' to use a V1 " - "Mamba cache manager.") if has_additional_snapshots and not use_v2: raise ValueError("Mamba additional snapshot offsets require " "use_kv_cache_manager_v2=True; V1 supports only " @@ -151,16 +167,29 @@ def get_kv_cache_manager_cls( # Skip Softmax only changes attention kernels. Hybrid models still # need a Mamba-capable cache manager for recurrent state. if is_disagg: - if kv_cache_config.enable_block_reuse: + backend, runtime = _resolve_disagg_transceiver_route( + cache_transceiver_config) + if use_v2: + if runtime != "PYTHON" or backend != "NIXL": + raise ValueError( + "KV cache manager V2 for hybrid Mamba disaggregated " + "serving requires transceiver_runtime='PYTHON' with " + "backend='NIXL'.") + else: + if (kv_cache_config.enable_block_reuse and runtime == "PYTHON"): + raise ValueError( + "Hybrid Mamba disaggregated serving with block reuse " + "and transceiver_runtime='PYTHON' requires " + "use_kv_cache_manager_v2=True.") + if kv_cache_config.enable_block_reuse: + return CppMambaHybridCacheManager + if runtime == "PYTHON" and backend == "NIXL": + logger.info("Python transceiver detected; using " + "MixedMambaHybridCacheManager for hybrid model") + return MixedMambaHybridCacheManager return CppMambaHybridCacheManager - if (cache_transceiver_config is not None and - cache_transceiver_config.transceiver_runtime == "PYTHON"): - logger.info("Python transceiver detected; using " - "MixedMambaHybridCacheManager for hybrid model") - return MixedMambaHybridCacheManager - return CppMambaHybridCacheManager - if use_py_mamba_cache_manager(): + if use_py_mamba_cache_manager() and not is_disagg: if use_v2: raise ValueError( "TRTLLM_USE_PY_MAMBA=1 conflicts with explicit " @@ -1913,6 +1942,8 @@ def _create_kv_cache_manager( manager_extra_kwargs = {} if issubclass(kv_cache_manager_cls, KVCacheManagerV2): manager_extra_kwargs["enable_stats"] = enable_kv_cache_stats + if issubclass(kv_cache_manager_cls, V2MambaHybridCacheManager): + manager_extra_kwargs["is_disagg"] = is_disagg if is_mla(config): kv_cache_manager = kv_cache_manager_cls( @@ -2016,7 +2047,9 @@ def _create_kv_cache_manager( and mamba_params.mamba_ssm_cache_dtype == torch.float16) mamba_manager_extra_kwargs = dict(manager_extra_kwargs) - if not issubclass(kv_cache_manager_cls, V2MambaHybridCacheManager): + if issubclass(kv_cache_manager_cls, V2MambaHybridCacheManager): + mamba_manager_extra_kwargs["conv_state_layout"] = "x_b_c" + else: mamba_manager_extra_kwargs["model_type"] = "nemotron_hybrid" kv_cache_manager = kv_cache_manager_cls( # mamba cache parameters @@ -2115,7 +2148,9 @@ def _create_kv_cache_manager( ("ENABLED" if use_replay else "DISABLED")) mamba_manager_extra_kwargs = dict(manager_extra_kwargs) - if not issubclass(kv_cache_manager_cls, V2MambaHybridCacheManager): + if issubclass(kv_cache_manager_cls, V2MambaHybridCacheManager): + mamba_manager_extra_kwargs["conv_state_layout"] = "q_k_v" + else: mamba_manager_extra_kwargs["model_type"] = "qwen3_next" kv_cache_manager = kv_cache_manager_cls( # mamba cache parameters diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index 348d47a9659c..9a2ed0ef52fd 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -16,7 +16,8 @@ from .llm_request import LlmRequest from .mamba_cache_manager import (BaseMambaCacheManager, CppMambaHybridCacheManager, - MixedMambaHybridCacheManager) + MixedMambaHybridCacheManager, + V2MambaHybridCacheManager) from .resource_manager import KVCacheManager CacheTransceiverCpp = tensorrt_llm.bindings.internal.batch_manager.CacheTransceiver @@ -157,14 +158,33 @@ def create_kv_cache_transceiver( "UCX_CUDA_IPC_ENABLE_MNNVL=n, UCX_RNDV_SCHEME=put_zcopy and/or unset UCX_NET_DEVICES upon server " "hangs or lower-than-expected performance.") - # Select transceiver implementation based on transceiver_runtime + # Select transceiver implementation based on transceiver_runtime. # transceiver_runtime == None or "CPP" -> use C++ transceiver (default) - # transceiver_runtime == "PYTHON" -> use Python transceiver - if cache_transceiver_config.transceiver_runtime == "PYTHON": - # Python transceiver currently only supports NIXL and DEFAULT backend - if cache_transceiver_config.backend not in ("DEFAULT", "NIXL"): + # transceiver_runtime == "PYTHON" -> use Python transceiver. + # + # V2MambaHybridCacheManager is backed by the Python KVCacheManagerV2 core, + # not the C++ BaseKVCacheManager binding required by CacheTransceiverCpp. + is_v2_mamba_hybrid = isinstance(mamba_cache_manager, + V2MambaHybridCacheManager) + use_python_transceiver = ( + cache_transceiver_config.transceiver_runtime == "PYTHON") + + if is_v2_mamba_hybrid and not use_python_transceiver: + raise ValueError( + "V2MambaHybridCacheManager requires transceiver_runtime='PYTHON' " + "with backend='NIXL'; it cannot use the C++ transceiver.") + + if use_python_transceiver: + if isinstance(mamba_cache_manager, CppMambaHybridCacheManager): + raise ValueError( + "transceiver_runtime='PYTHON' cannot drive " + "CppMambaHybridCacheManager (C++ pool backed). Use " + "transceiver_runtime='CPP', or select the V2 manager " + "with use_kv_cache_manager_v2=True.") + # DEFAULT has already been resolved above, so Python must see NIXL. + if cache_transceiver_config.backend != "NIXL": raise ValueError( - f"Python transceiver currently only supports NIXL or DEFAULT backend, " + f"Python transceiver currently only supports the NIXL backend, " f"got {cache_transceiver_config.backend}. " f"Please use transceiver_runtime='CPP' for MPI, UCX, or MOONCAKE backends." ) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 96550e5fe53e..b3de90d0d779 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -17,8 +17,8 @@ import os from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import (TYPE_CHECKING, Dict, Iterable, List, NamedTuple, Optional, - Tuple, Union) +from typing import (TYPE_CHECKING, Dict, Iterable, List, Literal, NamedTuple, + Optional, Tuple, Union) import torch import triton @@ -2466,8 +2466,12 @@ def __init__( is_draft: bool = False, use_replay_state_update: bool = False, mamba_ssm_stochastic_rounding: bool = False, + conv_state_layout: Literal["x_b_c", "q_k_v"] = "x_b_c", **kwargs, ) -> None: + if conv_state_layout not in ("x_b_c", "q_k_v"): + raise ValueError( + f"Unsupported convolution state layout: {conv_state_layout!r}") total_layers = len(mamba_layer_mask) if layer_mask is None: full_attention_layer_mask = [False] * total_layers @@ -2518,18 +2522,42 @@ def __init__( if self.local_num_mamba_layers > 0: tp_size = mapping.tp_size if not mapping.enable_attention_dp else 1 d_inner = mamba_head_dim * mamba_num_heads - conv_dim = d_inner + 2 * mamba_n_groups * mamba_d_state + grouped_state_dim = mamba_n_groups * mamba_d_state + conv_dim = d_inner + 2 * grouped_state_dim nheads = mamba_num_heads assert nheads % tp_size == 0, "mamba_num_heads must be divisible by tp_size" assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" + if kwargs.get("is_disagg", + False) and grouped_state_dim % tp_size != 0: + raise ValueError( + "Disaggregated Mamba transfer requires each convolution " + "state section to be divisible by tp_size") if use_replay_state_update: assert mamba_n_groups % tp_size == 0, \ "replay state update requires mamba_n_groups divisible by tp_size" self._n_groups_per_rank = mamba_n_groups // tp_size + d_inner_local = d_inner // tp_size + grouped_state_dim_local = grouped_state_dim // tp_size conv_dim = conv_dim // tp_size nheads = nheads // tp_size self.conv_state_shape = [conv_dim, mamba_d_conv - 1] self.ssm_state_shape = [nheads, mamba_head_dim, mamba_d_state] + # TP-mismatch disaggregated transfers must split the flat + # convolution state at its true semantic boundaries. Mamba2 stores + # [x | B | C], while GDN stores [Q | K | V]. The large section is + # therefore first for Mamba2 and last for GDN. + if conv_state_layout == "x_b_c": + self.conv_section_dims = [ + d_inner_local, + grouped_state_dim_local, + grouped_state_dim_local, + ] + else: + self.conv_section_dims = [ + grouped_state_dim_local, + grouped_state_dim_local, + d_inner_local, + ] self.ssm_count = math.prod(self.ssm_state_shape) self.conv_count = math.prod(self.conv_state_shape) self.ssm_bytes = self.ssm_count * self.ssm_state_dtype.itemsize @@ -2541,6 +2569,7 @@ def __init__( self._n_groups_per_rank = 0 self.conv_state_shape = [] self.ssm_state_shape = [] + self.conv_section_dims = [] self.ssm_count = 0 self.conv_count = 0 self.ssm_bytes = 0 diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 08ceedf61054..a8dd45606309 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -680,16 +680,6 @@ def drafting_loop_wrapper(model): config = model_engine.model.model_config.pretrained_config max_num_seq_slots = getattr(model_engine, "max_num_seq_slots", max_batch_size * getattr(mapping, "pp_size", 1)) - if is_hybrid_linear(config) and kv_cache_config.enable_block_reuse and ( - cache_transceiver_config is not None - and cache_transceiver_config.backend is not None - and cache_transceiver_config.transceiver_runtime == "PYTHON"): - logger.warning( - "Disabling block reuse for MambaHybridCacheManager-based models when disagg + Python transceiver enabled" - ) - kv_cache_config.enable_block_reuse = False - _set_model_engines_cache_reuse([model_engine, draft_model_engine], - False) if is_mla(config): if model_engine.model.model_config.enable_flash_mla: tokens_per_block = 64 @@ -907,17 +897,16 @@ def drafting_loop_wrapper(model): if is_disagg and is_hybrid: # NOTE: TRTLLM_USE_PY_MAMBA is an agg-mode-only override and has - # no effect in disagg. The disagg manager choice is driven solely - # by transceiver_runtime: PYTHON => PythonMambaCacheManager, - # otherwise CppMambaHybridCacheManager (unified pool, default). + # no effect in disagg. The disagg manager choice is driven by + # get_kv_cache_manager_cls and cache_transceiver_config. if os.environ.get("TRTLLM_USE_PY_MAMBA", "0") == "1": logger.warning( "TRTLLM_USE_PY_MAMBA is ignored in disaggregated serving; " - "use cache_transceiver_config.transceiver_runtime='PYTHON' " - "to select PythonMambaCacheManager.") + "configure transceiver_runtime='PYTHON' with backend='NIXL' " + "to select MixedMambaHybridCacheManager.") else: logger.info("Disaggregated serving with hybrid model detected. " - "Using CppMambaHybridCacheManager.") + "Using the configured Mamba cache manager.") # Get draft config for one-engine speculative decoding if available draft_config = getattr(model_engine.model, 'draft_config', None) diff --git a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py index 3d9a3f83dcc3..f5bdc6de75e4 100644 --- a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py +++ b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py @@ -23,6 +23,10 @@ from tensorrt_llm._torch.pyexecutor import py_executor as executor_module from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import BindKvCacheTransceiver from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( + CppMambaHybridCacheManager, + V2MambaHybridCacheManager, +) from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig @@ -519,6 +523,59 @@ def test_flag_unset_preserves_python_transceiver(monkeypatch): constructor.assert_called_once() +def test_python_nixl_transceiver_accepts_v2_mamba_manager(monkeypatch): + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime="PYTHON") + expected = object() + constructor = Mock(return_value=expected) + fake_module = SimpleNamespace(KvCacheTransceiverV2=constructor) + monkeypatch.setitem(sys.modules, "tensorrt_llm._torch.disaggregation.transceiver", fake_module) + manager = object.__new__(V2MambaHybridCacheManager) + + result = transceiver_module.create_kv_cache_transceiver( + Mock(), Mock(), manager, Mock(), config, manager + ) + + assert result is expected + constructor.assert_called_once() + + +@pytest.mark.parametrize("runtime", [None, "CPP", "auto"]) +def test_cpp_runtime_rejects_v2_mamba_manager(runtime): + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime=runtime) + manager = object.__new__(V2MambaHybridCacheManager) + + with pytest.raises(ValueError, match="requires transceiver_runtime='PYTHON'"): + transceiver_module.create_kv_cache_transceiver( + Mock(), Mock(), manager, Mock(), config, manager + ) + + +def test_python_runtime_rejects_cpp_mamba_manager(): + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime="PYTHON") + manager = object.__new__(CppMambaHybridCacheManager) + + with pytest.raises(ValueError, match="cannot drive CppMambaHybridCacheManager"): + transceiver_module.create_kv_cache_transceiver( + Mock(), Mock(), manager, Mock(), config, manager + ) + + +@pytest.mark.parametrize("runtime", [None, "CPP"]) +def test_cpp_runtime_keeps_cpp_mamba_manager(monkeypatch, runtime): + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime=runtime) + manager = object.__new__(CppMambaHybridCacheManager) + expected = object() + constructor = Mock(return_value=expected) + monkeypatch.setattr(transceiver_module, "BindKvCacheTransceiver", constructor) + + result = transceiver_module.create_kv_cache_transceiver( + Mock(), Mock(), manager, Mock(), config, manager + ) + + assert result is expected + constructor.assert_called_once() + + def test_flag_unset_preserves_libfabric_selection(monkeypatch): monkeypatch.setenv(transceiver_module._NIXL_KVCACHE_BACKEND_ENV, "LIBFABRIC") config = CacheTransceiverConfig(backend="NIXL") diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index e28046cd8bfb..db6f8aaa1b2d 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -9,6 +9,9 @@ import pytest import torch +from tensorrt_llm._torch.disaggregation.resource.kv_extractor import build_page_table_from_manager +from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup, MambaLayerGroup +from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata from tensorrt_llm._torch.pyexecutor._util import ( KvCacheCreator, @@ -46,12 +49,16 @@ from tensorrt_llm._utils import torch_dtype_to_binding from tensorrt_llm.bindings.internal.batch_manager import LinearCacheType from tensorrt_llm.llmapi.llm_args import ( + CacheTransceiverConfig, KvCacheConfig, MambaStateConfig, MTPDecodingConfig, TorchLlmArgs, ) -from tensorrt_llm.llmapi.llm_utils import _resolve_kv_cache_manager_v2_auto +from tensorrt_llm.llmapi.llm_utils import ( + _resolve_kv_cache_manager_v2_auto, + _resolve_transceiver_runtime_auto, +) from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import GpuCacheTierConfig, LayerId from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManager as RuntimeKVCacheManager @@ -240,6 +247,7 @@ def __init__(self, *args, **kwargs): assert captured_cpp["model_type"] == "qwen3_next" assert captured_v2["use_replay_state_update"] is False assert "model_type" not in captured_v2 + assert captured_v2["conv_state_layout"] == "q_k_v" def test_hybrid_cache_manager_factory_rejects_cpp_preference_with_explicit_v2( @@ -357,6 +365,93 @@ def test_hybrid_cache_manager_factory_routes_explicit_snapshots_to_v2( ) +@pytest.mark.parametrize("backend", ["NIXL", "DEFAULT"]) +def test_hybrid_cache_manager_factory_routes_explicit_v2_disagg(monkeypatch, backend): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + for env_var in ( + "TRTLLM_USE_NIXL_KVCACHE", + "TRTLLM_USE_UCX_KVCACHE", + "TRTLLM_USE_MOONCAKE_KVCACHE", + "TRTLLM_USE_MPI_KVCACHE", + ): + monkeypatch.delenv(env_var, raising=False) + assert ( + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + use_kv_cache_manager_v2=True, + ), + is_disagg=True, + cache_transceiver_config=CacheTransceiverConfig( + backend=backend, transceiver_runtime="PYTHON" + ), + ) + is V2MambaHybridCacheManager + ) + + +def test_hybrid_cache_manager_factory_rejects_python_v1_disagg_reuse(monkeypatch): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + + with pytest.raises(ValueError, match="requires use_kv_cache_manager_v2=True"): + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + use_kv_cache_manager_v2=False, + ), + is_disagg=True, + cache_transceiver_config=CacheTransceiverConfig( + backend="NIXL", + transceiver_runtime="PYTHON", + ), + ) + + +@pytest.mark.parametrize( + ("backend", "runtime", "backend_env"), + [ + ("DEFAULT", "PYTHON", "TRTLLM_USE_UCX_KVCACHE"), + ("UCX", "PYTHON", None), + ("NIXL", "auto", None), + ("NIXL", None, None), + ("NIXL", "CPP", None), + ("UCX", None, None), + ], +) +def test_hybrid_cache_manager_factory_rejects_unsupported_v2_disagg_route( + monkeypatch, backend, runtime, backend_env +): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + for env_var in ( + "TRTLLM_USE_NIXL_KVCACHE", + "TRTLLM_USE_UCX_KVCACHE", + "TRTLLM_USE_MOONCAKE_KVCACHE", + "TRTLLM_USE_MPI_KVCACHE", + ): + monkeypatch.delenv(env_var, raising=False) + if backend_env is not None: + monkeypatch.setenv(backend_env, "1") + + with pytest.raises(ValueError, match="requires transceiver_runtime='PYTHON'"): + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + use_kv_cache_manager_v2=True, + ), + is_disagg=True, + cache_transceiver_config=CacheTransceiverConfig( + backend=backend, transceiver_runtime=runtime + ), + ) + + @pytest.mark.parametrize( ("env_name", "env_value", "expected_error"), [ @@ -413,35 +508,52 @@ def test_hybrid_cache_manager_factory_keeps_v1_disagg_route(monkeypatch, use_v2) use_kv_cache_manager_v2=use_v2, ), is_disagg=True, - cache_transceiver_config=SimpleNamespace(transceiver_runtime="PYTHON"), + cache_transceiver_config=CacheTransceiverConfig( + backend="NIXL", transceiver_runtime="PYTHON" + ), ) is MixedMambaHybridCacheManager ) -@pytest.mark.parametrize("enable_block_reuse", [False, True]) -@pytest.mark.parametrize("transceiver_runtime", [None, "PYTHON"]) -def test_hybrid_cache_manager_factory_rejects_v2_disagg( - monkeypatch, enable_block_reuse, transceiver_runtime -): - monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) - monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) - transceiver_config = ( - None - if transceiver_runtime is None - else SimpleNamespace(transceiver_runtime=transceiver_runtime) - ) +def test_hybrid_models_resolve_auto_to_python_transceiver(monkeypatch): + from tensorrt_llm._torch.models.modeling_nemotron_h import NemotronHForCausalLM + from tensorrt_llm._torch.models.modeling_qwen3_5 import Qwen3_5VLModel + from tensorrt_llm._torch.models.modeling_qwen3_next import Qwen3NextForCausalLM - with pytest.raises(ValueError, match="V2.*not supported with disaggregated"): - get_kv_cache_manager_cls( - _hybrid_model_config(), - KvCacheConfig( - enable_block_reuse=enable_block_reuse, - use_kv_cache_manager_v2=True, - ), - is_disagg=True, - cache_transceiver_config=transceiver_config, + for env_var in ( + "TRTLLM_USE_NIXL_KVCACHE", + "TRTLLM_USE_UCX_KVCACHE", + "TRTLLM_USE_MOONCAKE_KVCACHE", + "TRTLLM_USE_MPI_KVCACHE", + ): + monkeypatch.delenv(env_var, raising=False) + + for model_cls in (NemotronHForCausalLM, Qwen3NextForCausalLM, Qwen3_5VLModel): + llm_args = TorchLlmArgs( + model="/tmp/dummy_model", + cache_transceiver_config=CacheTransceiverConfig(backend="DEFAULT"), ) + _resolve_transceiver_runtime_auto(llm_args, model_cls) + assert llm_args.cache_transceiver_config.transceiver_runtime == "PYTHON" + + +def test_v2_disagg_slice_skips_state_index_on_mamba_free_pp_rank(): + manager = object.__new__(V2MambaHybridCacheManager) + manager.local_num_mamba_layers = 0 + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._kv_cache_manager = manager + transceiver._reuse_adapter = SimpleNamespace(tokens_per_block=32) + transceiver._page_table = SimpleNamespace(layer_groups=[]) + request = SimpleNamespace( + is_generation_only_request=lambda: False, + prompt_len=0, + py_request_id=123, + ) + + kv_slice = transceiver._create_kv_slice(request) + + assert kv_slice.mamba_state_index is None @pytest.mark.parametrize( @@ -1453,6 +1565,7 @@ def _build_v2_hybrid_with_mamba_layer( enable_attention_dp=False, enable_swa_scratch_reuse=False, dtype=DataType.HALF, + conv_state_layout="x_b_c", ): """Construct a real V2MambaHybridCacheManager.""" mamba_mask = [True] * num_mamba_layers + [False] * num_attention_layers @@ -1496,6 +1609,7 @@ def _build_v2_hybrid_with_mamba_layer( vocab_size=1024, use_replay_state_update=use_replay_state_update, dtype=dtype, + conv_state_layout=conv_state_layout, ) @@ -1841,6 +1955,47 @@ def test_v2_hybrid_swa_scratch_keeps_ssm_placeholder_rows(): mgr.shutdown() +@skip_no_cuda +def test_v2_hybrid_disagg_page_table_preserves_lifecycle_indices(): + mgr = _build_v2_hybrid_with_mamba_layer(max_batch_size=4, num_mamba_layers=2) + try: + page_table = build_page_table_from_manager(mgr) + + assert len(page_table.layer_groups) == mgr.impl._storage.num_life_cycles + assert isinstance(page_table.layer_groups[0], MambaLayerGroup) + assert isinstance(page_table.layer_groups[1], AttentionLayerGroup) + + requests = mgr.add_dummy_requests([123], token_nums=[64], is_gen=False) + assert len(requests) == 1 + attention_blocks = list( + mgr.kv_cache_map[123].get_aggregated_page_indices(1, valid_only=True) + ) + assert attention_blocks + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_disagg_page_table_uses_qwen3_next_conv_sections(): + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=4, + conv_state_layout="q_k_v", + ) + try: + page_table = build_page_table_from_manager(mgr) + mamba_group = page_table.layer_groups[0] + + assert isinstance(mamba_group, MambaLayerGroup) + d_conv_m1 = mgr.conv_state_shape[1] + conv_elem_size = mgr.all_conv_states[0].element_size() + assert mamba_group.conv_section_bytes == [ + dim * d_conv_m1 * conv_elem_size for dim in mgr.conv_section_dims + ] + assert mgr.conv_section_dims == [8, 8, 32] + finally: + mgr.shutdown() + + @skip_no_cuda def test_cpp_hybrid_replay_buffers_size_by_tokens_per_gen_step(): spec_config = _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5) diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index 6fc2f68fef68..e31180ce18ec 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -1,3 +1,17 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import numpy as np import pytest @@ -13,6 +27,7 @@ get_num_layer_groups, get_num_layers, get_physical_pool, + get_slot_address, get_unique_layers, ) from tensorrt_llm._torch.pyexecutor.resource_manager import ( @@ -235,8 +250,18 @@ def test_layer_group_meta_serialization(): def test_mamba_layer_group_serialization(): from tensorrt_llm._torch.disaggregation.resource.page import MambaLayerGroup, PhysicalPool - conv_pool = PhysicalPool(base_address=1000, slot_bytes=128, num_slots=10) - ssm_pool = PhysicalPool(base_address=8000, slot_bytes=256, num_slots=8) + conv_pool = PhysicalPool( + base_address=1000, + slot_bytes=128, + num_slots=10, + slot_stride_bytes=512, + ) + ssm_pool = PhysicalPool( + base_address=8000, + slot_bytes=256, + num_slots=8, + slot_stride_bytes=1024, + ) mlg = MambaLayerGroup( pool_group_idx=1, mamba_layer_offsets={10: 0, 11: 1, 12: 2}, @@ -244,11 +269,17 @@ def test_mamba_layer_group_serialization(): ssm_states=ssm_pool, conv_section_bytes=[512, 256, 256], ssm_bytes_per_head=128, + conv_layer_slot0_addresses={10: 1000, 11: 2000, 12: 3000}, + ssm_layer_slot0_addresses={10: 8000, 11: 9000, 12: 10000}, ) d = mlg.to_dict() assert d["mamba_layer_offsets"] == {10: 0, 11: 1, 12: 2} assert d["conv_section_bytes"] == [512, 256, 256] + assert d["conv_layer_slot0_addresses"] == {10: 1000, 11: 2000, 12: 3000} + assert d["ssm_layer_slot0_addresses"] == {10: 8000, 11: 9000, 12: 10000} + assert d["conv_states"]["slot_stride_bytes"] == 512 + assert d["ssm_states"]["slot_stride_bytes"] == 1024 from tensorrt_llm._torch.disaggregation.resource.page import LayerGroup @@ -258,11 +289,102 @@ def test_mamba_layer_group_serialization(): assert restored.conv_states.base_address == 1000 assert restored.conv_states.slot_bytes == 128 assert restored.conv_states.num_slots == 10 + assert restored.conv_states.slot_stride_bytes == 512 + assert get_slot_address(restored.conv_states, 3) == 1000 + 3 * 512 assert restored.ssm_states.base_address == 8000 assert restored.ssm_states.slot_bytes == 256 assert restored.ssm_states.num_slots == 8 + assert restored.ssm_states.slot_stride_bytes == 1024 assert restored.conv_section_bytes == [512, 256, 256] assert restored.ssm_bytes_per_head == 128 + assert restored.conv_layer_slot0_addresses == {10: 1000, 11: 2000, 12: 3000} + assert restored.ssm_layer_slot0_addresses == {10: 8000, 11: 9000, 12: 10000} + + legacy_pool = PhysicalPool.from_dict({"base_address": 1000, "slot_bytes": 128, "num_slots": 10}) + assert legacy_pool.slot_stride_bytes == legacy_pool.slot_bytes + + +def test_v2_mamba_registration_uses_coalesced_physical_pool(): + from tensorrt_llm._torch.disaggregation.resource.page import ( + KVCachePageTable, + MambaLayerGroup, + PhysicalPool, + PhysicalPoolGroup, + ) + from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs + + state_bytes = 64 + num_layers = 2 + num_slots = 8 + # Equal-sized SSM and convolution states share one interleaved V2 pool. + physical_slot_bytes = state_bytes * num_layers * 2 + physical_pool = PhysicalPool( + base_address=1000, + slot_bytes=physical_slot_bytes, + num_slots=num_slots, + ) + mamba_group = MambaLayerGroup( + pool_group_idx=0, + mamba_layer_offsets={1: 0, 2: 1}, + conv_states=PhysicalPool( + base_address=1000 + state_bytes, + slot_bytes=state_bytes, + num_slots=num_slots, + slot_stride_bytes=physical_slot_bytes, + ), + ssm_states=PhysicalPool( + base_address=1000, + slot_bytes=state_bytes, + num_slots=num_slots, + slot_stride_bytes=physical_slot_bytes, + ), + conv_layer_slot0_addresses={ + 1: 1000 + state_bytes, + 2: 1000 + state_bytes * 3, + }, + ssm_layer_slot0_addresses={ + 1: 1000, + 2: 1000 + state_bytes * 2, + }, + ) + page_table = KVCachePageTable( + tokens_per_block=16, + layer_groups=[mamba_group], + pool_groups=[PhysicalPoolGroup(pools=[physical_pool])], + ) + + assert get_unique_pool_memory_descs(page_table, device_id=3) == [ + (1000, physical_slot_bytes * num_slots, 3, "kv_cache_memory_pool0") + ] + + +def test_legacy_mamba_registration_uses_layer_major_pools(): + from tensorrt_llm._torch.disaggregation.resource.page import ( + KVCachePageTable, + MambaLayerGroup, + PhysicalPool, + ) + from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs + + num_layers = 3 + conv_pool = PhysicalPool(base_address=1000, slot_bytes=128, num_slots=10) + ssm_pool = PhysicalPool(base_address=8000, slot_bytes=256, num_slots=8) + mamba_group = MambaLayerGroup( + pool_group_idx=0, + mamba_layer_offsets={10: 0, 11: 1, 12: 2}, + conv_states=conv_pool, + ssm_states=ssm_pool, + ) + page_table = KVCachePageTable( + tokens_per_block=16, + layer_groups=[mamba_group], + pool_groups=[], + ) + + assert get_unique_pool_memory_descs(page_table, device_id=3) == [ + (1000, num_layers * conv_pool.num_slots * conv_pool.slot_bytes, 3, "kv_cache_memory_pool0"), + (8000, num_layers * ssm_pool.num_slots * ssm_pool.slot_bytes, 3, "kv_cache_memory_pool1"), + ] def test_mixed_page_table_serialization(): diff --git a/tests/unittest/disaggregated/test_mamba_transfer.py b/tests/unittest/disaggregated/test_mamba_transfer.py index 0cf0b3209899..0b5c3716c95c 100644 --- a/tests/unittest/disaggregated/test_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_mamba_transfer.py @@ -16,6 +16,7 @@ import uuid from typing import Dict, List +import numpy as np import pytest import torch @@ -23,13 +24,19 @@ import tensorrt_llm.bindings import tensorrt_llm.tensorrt_llm_transfer_agent_binding # noqa: F401 from tensorrt_llm import DisaggregatedParams, Mapping, SamplingParams +from tensorrt_llm._torch.disaggregation.native import rank_info +from tensorrt_llm._torch.disaggregation.native.mixers.ssm import peer +from tensorrt_llm._torch.disaggregation.resource import page from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 from tensorrt_llm._torch.pyexecutor.llm_request import ( ATTENTION_DP_DUMMY_REQUEST_ID, LlmRequest, LlmRequestType, ) -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MixedMambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( + MixedMambaHybridCacheManager, + V2MambaHybridCacheManager, +) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp @@ -183,8 +190,14 @@ def _init(rank): return results -def _create_managers(tp, max_batch_size=MAX_BATCH_SIZE, enable_attention_dp=False): - """Create MixedMambaHybridCacheManagers for all TP ranks (PP=1). +def _create_managers( + tp, + max_batch_size=MAX_BATCH_SIZE, + enable_attention_dp=False, + use_v2=False, + conv_state_layout="x_b_c", +): + """Create Mamba hybrid cache managers for all TP ranks (PP=1). Layer 0 is a dummy attention layer required by page table infrastructure. Layers 1..NUM_MAMBA_LAYERS are mamba layers under test. @@ -194,7 +207,16 @@ def _create_managers(tp, max_batch_size=MAX_BATCH_SIZE, enable_attention_dp=Fals mapping = Mapping( world_size=tp, rank=rank, tp_size=tp, pp_size=1, enable_attention_dp=enable_attention_dp ) - mgr = MixedMambaHybridCacheManager( + manager_cls = V2MambaHybridCacheManager if use_v2 else MixedMambaHybridCacheManager + manager_kwargs = ( + { + "is_disagg": True, + "conv_state_layout": conv_state_layout, + } + if use_v2 + else {} + ) + mgr = manager_cls( mamba_d_state=MAMBA_D_STATE, mamba_d_conv=MAMBA_D_CONV, mamba_num_heads=MAMBA_NUM_HEADS, @@ -220,19 +242,169 @@ def _create_managers(tp, max_batch_size=MAX_BATCH_SIZE, enable_attention_dp=Fals max_batch_size=max_batch_size, mapping=mapping, dtype=DataType.FLOAT, + **manager_kwargs, ) managers.append(mgr) return managers +def _mamba_layer_ids(manager): + if isinstance(manager, V2MambaHybridCacheManager): + return manager.mamba_layer_offsets + return manager._impl.mamba_layer_offsets + + +def _mamba_state_slot(manager, request_id): + if isinstance(manager, V2MambaHybridCacheManager): + return manager.get_state_indices([request_id])[0] + return manager.mamba_cache_index[request_id] + + +def _zero_mamba_states(manager): + for layer_idx in _mamba_layer_ids(manager): + manager.get_conv_states(layer_idx).zero_() + manager.get_ssm_states(layer_idx).zero_() + + +def test_mamba_policy_layer_major_v1_ptrs(): + pool = page.PhysicalPool(base_address=100, slot_bytes=10, num_slots=8) + + ptrs = peer.MambaPolicy._build_layer_ptrs( + pool=pool, + layer_offsets={1: 0, 2: 1}, + overlapping_layers=[1, 2], + slot=3, + ) + + np.testing.assert_array_equal(ptrs, [130, 210]) + + +def test_mamba_policy_slot_major_layer_ptrs(): + """V2 Mamba state tensors step by physical slots, not V1 layers.""" + self_mlg = page.MambaLayerGroup( + pool_group_idx=0, + mamba_layer_offsets={1: 0, 2: 1}, + conv_states=page.PhysicalPool( + base_address=1000, + slot_bytes=10, + num_slots=8, + slot_stride_bytes=20, + ), + ssm_states=page.PhysicalPool( + base_address=3000, + slot_bytes=20, + num_slots=8, + slot_stride_bytes=40, + ), + conv_section_bytes=[10], + ssm_bytes_per_head=10, + conv_layer_slot0_addresses={1: 1000, 2: 1010}, + ssm_layer_slot0_addresses={1: 3000, 2: 3020}, + ) + peer_mlg = page.MambaLayerGroup( + pool_group_idx=0, + mamba_layer_offsets={1: 0, 2: 1}, + conv_states=page.PhysicalPool( + base_address=5000, + slot_bytes=10, + num_slots=8, + slot_stride_bytes=20, + ), + ssm_states=page.PhysicalPool( + base_address=7000, + slot_bytes=20, + num_slots=8, + slot_stride_bytes=40, + ), + conv_section_bytes=[10], + ssm_bytes_per_head=10, + conv_layer_slot0_addresses={1: 5000, 2: 5010}, + ssm_layer_slot0_addresses={1: 7000, 2: 7020}, + ) + self_ri = rank_info.RankInfo( + instance_name="self", + instance_rank=0, + tp_size=1, + tp_rank=0, + pp_size=1, + pp_rank=0, + layer_num_per_pp=[2], + sender_endpoints=[], + server_endpoint="", + self_endpoint="", + transfer_engine_info=bytes(), + ) + peer_ri = rank_info.RankInfo( + instance_name="peer", + instance_rank=0, + tp_size=1, + tp_rank=0, + pp_size=1, + pp_rank=0, + layer_num_per_pp=[2], + sender_endpoints=[], + server_endpoint="", + self_endpoint="", + transfer_engine_info=bytes(), + ) + + src_frags, dst_frags, kv_sizes = peer.MambaPolicy.build_mamba_frags( + self_mlg=self_mlg, + peer_mlg=peer_mlg, + src_slot=3, + dst_slot=5, + self_ri=self_ri, + peer_ri=peer_ri, + ) + + assert src_frags == [1060, 1070, 3120, 3140] + assert dst_frags == [5100, 5110, 7200, 7220] + assert kv_sizes == [10, 10, 20, 20] + + +def test_mamba_policy_slot_major_interleaved_role_ptrs(): + """Layer/role offsets remain inside each coalesced V2 physical slot.""" + state_bytes = 64 + physical_slot_bytes = 4 * state_bytes + conv_pool = page.PhysicalPool( + base_address=1000 + state_bytes, + slot_bytes=state_bytes, + num_slots=8, + slot_stride_bytes=physical_slot_bytes, + ) + + ptrs = peer.MambaPolicy._build_layer_ptrs( + pool=conv_pool, + layer_offsets={1: 0, 2: 1}, + overlapping_layers=[1, 2], + slot=3, + layer_slot0_addresses={ + 1: 1000 + state_bytes, + 2: 1000 + 3 * state_bytes, + }, + ) + + np.testing.assert_array_equal( + ptrs, + [ + 1000 + state_bytes + 3 * physical_slot_bytes, + 1000 + 3 * state_bytes + 3 * physical_slot_bytes, + ], + ) + + # --------------------------------------------------------------------------- # Ground truth: generate, shard, write, compute expected, read actual # --------------------------------------------------------------------------- -def _full_conv_section_dims() -> List[int]: - """Full (unsharded) first-dim sizes: [x(d_inner) | B(ng*ds) | C(ng*ds)].""" +def _full_conv_section_dims(conv_state_layout="x_b_c") -> List[int]: + """Full first-dimension sizes in the model's convolution-state order.""" d_inner = MAMBA_HEAD_DIM * MAMBA_NUM_HEADS ng_ds = MAMBA_N_GROUPS * MAMBA_D_STATE - return [d_inner, ng_ds, ng_ds] + if conv_state_layout == "x_b_c": + return [d_inner, ng_ds, ng_ds] + if conv_state_layout == "q_k_v": + return [ng_ds, ng_ds, d_inner] + raise ValueError(f"Unsupported convolution state layout: {conv_state_layout!r}") def _generate_ground_truth(num_requests: int, seed: int = 12345): @@ -270,29 +442,48 @@ def _shard_ssm(full_ssm: torch.Tensor, tp: int, tp_rank: int) -> torch.Tensor: return full_ssm[tp_rank * n : (tp_rank + 1) * n].clone() -def _shard_conv(full_conv: torch.Tensor, tp: int, tp_rank: int) -> torch.Tensor: - """Shard conv per-section along dim 0: [x | B | C] each independently.""" +def _shard_conv( + full_conv: torch.Tensor, + tp: int, + tp_rank: int, + conv_state_layout="x_b_c", +) -> torch.Tensor: + """Shard each semantic convolution-state section independently.""" parts = [] offset = 0 - for sec_dim in _full_conv_section_dims(): + for sec_dim in _full_conv_section_dims(conv_state_layout): n = sec_dim // tp parts.append(full_conv[offset + tp_rank * n : offset + (tp_rank + 1) * n]) offset += sec_dim return torch.cat(parts, dim=0).clone() -def _write_ground_truth_to_ctx(managers, tp, ground_truth, request_ids): +def _write_ground_truth_to_ctx( + managers, + tp, + ground_truth, + request_ids, + conv_state_layout="x_b_c", +): """Write sharded ground truth into ctx managers' allocated mamba slots.""" for rank, mgr in enumerate(managers): for req_idx, rid in enumerate(request_ids): - slot = mgr.mamba_cache_index[rid] - for layer_idx in mgr._impl.mamba_layer_offsets: + slot = _mamba_state_slot(mgr, rid) + for layer_idx in _mamba_layer_ids(mgr): full = ground_truth[req_idx][layer_idx] mgr.get_ssm_states(layer_idx)[slot] = _shard_ssm(full["ssm"], tp, rank) - mgr.get_conv_states(layer_idx)[slot] = _shard_conv(full["conv"], tp, rank) - - -def _compute_expected(ground_truth, gen_managers, gen_tp, gen_request_ids) -> Dict: + mgr.get_conv_states(layer_idx)[slot] = _shard_conv( + full["conv"], tp, rank, conv_state_layout + ) + + +def _compute_expected( + ground_truth, + gen_managers, + gen_tp, + gen_request_ids, + conv_state_layout="x_b_c", +) -> Dict: """Compute expected mamba states BEFORE transfer. Returns: {(gen_rank, req_idx, layer_idx): {"conv": Tensor, "ssm": Tensor}} @@ -300,11 +491,11 @@ def _compute_expected(ground_truth, gen_managers, gen_tp, gen_request_ids) -> Di expected = {} for gen_rank, mgr in enumerate(gen_managers): for req_idx in range(len(gen_request_ids)): - for layer_idx in mgr._impl.mamba_layer_offsets: + for layer_idx in _mamba_layer_ids(mgr): full = ground_truth[req_idx][layer_idx] expected[(gen_rank, req_idx, layer_idx)] = { "ssm": _shard_ssm(full["ssm"], gen_tp, gen_rank), - "conv": _shard_conv(full["conv"], gen_tp, gen_rank), + "conv": _shard_conv(full["conv"], gen_tp, gen_rank, conv_state_layout), } return expected @@ -317,8 +508,8 @@ def _read_actual(gen_managers, gen_request_ids) -> Dict: actual = {} for gen_rank, mgr in enumerate(gen_managers): for req_idx, rid in enumerate(gen_request_ids): - slot = mgr.mamba_cache_index[rid] - for layer_idx in mgr._impl.mamba_layer_offsets: + slot = _mamba_state_slot(mgr, rid) + for layer_idx in _mamba_layer_ids(mgr): actual[(gen_rank, req_idx, layer_idx)] = { "conv": mgr.get_conv_states(layer_idx)[slot].cpu().clone(), "ssm": mgr.get_ssm_states(layer_idx)[slot].cpu().clone(), @@ -358,14 +549,26 @@ def test_mamba_disagg_attention_dp_dummy_with_batch_size_one(): mgr.shutdown() -def run_mamba_transfer_test(ctx_tp: int, gen_tp: int): +def run_mamba_transfer_test( + ctx_tp: int, + gen_tp: int, + use_v2: bool = False, + conv_state_layout: str = "x_b_c", +): """Test mamba transfer: ctx_tp -> gen_tp (PP=1, no DP).""" # -- 1. Create managers, zero mamba caches -- - ctx_mgrs = _create_managers(ctx_tp) - gen_mgrs = _create_managers(gen_tp) + ctx_mgrs = _create_managers( + ctx_tp, + use_v2=use_v2, + conv_state_layout=conv_state_layout, + ) + gen_mgrs = _create_managers( + gen_tp, + use_v2=use_v2, + conv_state_layout=conv_state_layout, + ) for mgr in ctx_mgrs + gen_mgrs: - mgr._impl.mamba_cache.conv.zero_() - mgr._impl.mamba_cache.temporal.zero_() + _zero_mamba_states(mgr) # -- 2. Create transceivers -- config = CacheTransceiverConfig( @@ -419,14 +622,29 @@ def run_mamba_transfer_test(ctx_tp: int, gen_tp: int): # -- 4. Allocate slots -- ctx_batch = ScheduledRequests() - ctx_batch.reset_context_requests(ctx_reqs) - for mgr in ctx_mgrs: - mgr.prepare_resources(ctx_batch) + if use_v2: + ctx_batch.context_requests_last_chunk = ctx_reqs + for mgr in ctx_mgrs: + for req in ctx_reqs: + assert mgr.prepare_context(req) + assert mgr.resize_context(req, req.context_chunk_size) + mgr.prepare_resources(ctx_batch) + else: + ctx_batch.reset_context_requests(ctx_reqs) + for mgr in ctx_mgrs: + mgr.prepare_resources(ctx_batch) gen_batch = ScheduledRequests() - gen_batch.reset_context_requests(gen_reqs) - for mgr in gen_mgrs: - mgr.prepare_resources(gen_batch) + if use_v2: + gen_batch.context_requests_last_chunk = gen_reqs + for mgr in gen_mgrs: + for req in gen_reqs: + assert mgr.prepare_disagg_gen_init(req) + mgr.prepare_resources(gen_batch) + else: + gen_batch.reset_context_requests(gen_reqs) + for mgr in gen_mgrs: + mgr.prepare_resources(gen_batch) for req in ctx_reqs + gen_reqs: req.context_current_position = req.prompt_len @@ -438,10 +656,22 @@ def run_mamba_transfer_test(ctx_tp: int, gen_tp: int): # -- 5. Ground truth -> shard -> write to ctx -- ground_truth = _generate_ground_truth(len(REQUEST_LENGTHS)) - _write_ground_truth_to_ctx(ctx_mgrs, ctx_tp, ground_truth, ctx_rids) + _write_ground_truth_to_ctx( + ctx_mgrs, + ctx_tp, + ground_truth, + ctx_rids, + conv_state_layout, + ) # -- 6. Compute expected BEFORE transfer -- - expected = _compute_expected(ground_truth, gen_mgrs, gen_tp, gen_rids) + expected = _compute_expected( + ground_truth, + gen_mgrs, + gen_tp, + gen_rids, + conv_state_layout, + ) # -- 7. Transfer -- for rank in range(gen_tp): @@ -499,3 +729,23 @@ def test_mamba_transfer(ctx_tp, gen_tp): print(f"\nMamba transfer test: ctx_tp={ctx_tp} -> gen_tp={gen_tp}") run_mamba_transfer_test(ctx_tp, gen_tp) print("PASSED") + + +@pytest.mark.timeout(180) +@pytest.mark.parametrize( + "ctx_tp,gen_tp,conv_state_layout", + [ + (2, 2, "x_b_c"), + (2, 4, "q_k_v"), + (4, 2, "x_b_c"), + ], + ids=["same_tp_xbc", "expand_tp_qkv", "contract_tp_xbc"], +) +def test_v2_mamba_transfer(ctx_tp, gen_tp, conv_state_layout): + """Transfer slot-major V2 Mamba states through Python/NIXL.""" + run_mamba_transfer_test( + ctx_tp, + gen_tp, + use_v2=True, + conv_state_layout=conv_state_layout, + ) From 66558d5794ef383d85981e35cca1abc947c211fa Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 21 Jul 2026 20:43:03 +0800 Subject: [PATCH 05/22] [TRTLLM-11875][refactor] Derive Mamba state addresses from strides (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../disaggregation/native/mixers/ssm/peer.py | 40 ++------ .../disaggregation/resource/kv_extractor.py | 92 ++++++++++++------- .../_torch/disaggregation/resource/page.py | 44 +++------ .../_torch/disaggregation/resource/utils.py | 10 +- .../disaggregated/region/test_page.py | 12 ++- .../unittest/disaggregated/test_extractor.py | 59 +++++++++--- .../disaggregated/test_mamba_transfer.py | 13 +-- 7 files changed, 148 insertions(+), 122 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py index 81feb03a25a7..a96fc3e8f496 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py @@ -351,28 +351,18 @@ def _build_layer_ptrs( layer_offsets: Dict[int, int], overlapping_layers: List[int], slot: int, - layer_slot0_addresses: Optional[Dict[int, int]] = None, ) -> np.ndarray: - """Build per-layer pointers for a given pool (conv or SSM) and slot. - - V1 stores states layer-major, so its layer base is derived from the - Mamba-local layer offset. V2 stores buffers inside coalesced slot-major - pools; its manager-provided slot-0 addresses preserve the per-layer - offsets within that physical slot. - """ - ptrs = [] + """Build per-layer pointers from a pool's affine layer/slot layout.""" slot_stride_bytes = pool.slot_stride_bytes + layer_stride_bytes = pool.layer_stride_bytes assert slot_stride_bytes is not None - for glid in overlapping_layers: - if layer_slot0_addresses is not None: - ptrs.append(layer_slot0_addresses[glid] + slot * slot_stride_bytes) - else: - lid = layer_offsets[glid] - ptrs.append( - pool.base_address - + lid * pool.num_slots * pool.slot_bytes - + slot * pool.slot_bytes - ) + assert layer_stride_bytes is not None + ptrs = [ + pool.base_address + + layer_offsets[global_layer_id] * layer_stride_bytes + + slot * slot_stride_bytes + for global_layer_id in overlapping_layers + ] return np.array(ptrs, dtype=np.int64) @staticmethod @@ -458,29 +448,17 @@ def build_mamba_frags( (self_mlg.conv_states, peer_mlg.conv_states, True), (self_mlg.ssm_states, peer_mlg.ssm_states, False), ]: - self_layer_slot0_addresses = ( - self_mlg.conv_layer_slot0_addresses - if is_conv - else self_mlg.ssm_layer_slot0_addresses - ) - peer_layer_slot0_addresses = ( - peer_mlg.conv_layer_slot0_addresses - if is_conv - else peer_mlg.ssm_layer_slot0_addresses - ) src_ptrs = MambaPolicy._build_layer_ptrs( self_pool, self_mlg.mamba_layer_offsets, overlapping_layers, src_slot, - self_layer_slot0_addresses, ) dst_ptrs = MambaPolicy._build_layer_ptrs( peer_pool, peer_mlg.mamba_layer_offsets, overlapping_layers, dst_slot, - peer_layer_slot0_addresses, ) src_region = SpecRegion( diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index ddf876bd7dbd..599fd06bd560 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -12,9 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Dict, List +from typing import Dict, List, Sequence import numpy as np +import torch from tensorrt_llm._torch.disaggregation.base.region import ( DataLayout, @@ -120,12 +121,14 @@ def _build_layer_group_for_mamba( base_address=conv_state.data_ptr(), slot_bytes=conv_state.stride(1) * conv_state.element_size(), num_slots=conv_state.shape[1], + layer_stride_bytes=conv_state.stride(0) * conv_state.element_size(), ) ssm_pool = PhysicalPool( base_address=ssm_state.data_ptr(), slot_bytes=ssm_state.stride(1) * ssm_state.element_size(), num_slots=ssm_state.shape[1], + layer_stride_bytes=ssm_state.stride(0) * ssm_state.element_size(), ) # Per-section bytes for conv_state and per-head bytes for ssm_state. @@ -156,6 +159,46 @@ def _slot_stride_bytes(tensor) -> int: return int(tensor.stride(0) * tensor.element_size()) +def _build_v2_mamba_state_pool(states: Sequence[torch.Tensor]) -> PhysicalPool: + """Describe affine layer/slot addressing for one V2 Mamba state role.""" + if not states: + raise ValueError("V2 Mamba state pool requires at least one layer") + + first_state = states[0] + base_address = int(first_state.data_ptr()) + num_slots = int(first_state.shape[0]) + slot_bytes = int(first_state[0].numel() * first_state.element_size()) + slot_stride_bytes = _slot_stride_bytes(first_state) + + num_layers = len(states) + if slot_stride_bytes % num_layers != 0: + raise ValueError("V2 Mamba physical slot must divide evenly across layers") + # Each role appears once per layer in its size-class pool. Equal-size SSM + # and convolution states share that pool and are interleaved, so their + # layer stride includes both role payloads. + layer_stride_bytes = slot_stride_bytes // num_layers + + for layer_offset, state in enumerate(states): + state_slot_bytes = int(state[0].numel() * state.element_size()) + if ( + int(state.shape[0]) != num_slots + or state_slot_bytes != slot_bytes + or _slot_stride_bytes(state) != slot_stride_bytes + ): + raise ValueError("V2 Mamba state tensors must share one slot layout per role") + expected_address = base_address + layer_offset * layer_stride_bytes + if int(state.data_ptr()) != expected_address: + raise ValueError("V2 Mamba state tensors must have a uniform layer stride per role") + + return PhysicalPool( + base_address=base_address, + slot_bytes=slot_bytes, + num_slots=num_slots, + slot_stride_bytes=slot_stride_bytes, + layer_stride_bytes=layer_stride_bytes, + ) + + def _build_layer_group_for_v2_mamba( manager: V2MambaHybridCacheManager, pool_group_idx: int ) -> MambaLayerGroup: @@ -164,27 +207,20 @@ def _build_layer_group_for_v2_mamba( for global_layer_id, local_layer_id in manager.mamba_layer_offsets.items() } + expected_offsets = list(range(len(mamba_layer_offsets))) + if sorted(mamba_layer_offsets.values()) != expected_offsets: + raise ValueError("V2 Mamba layer offsets must be dense") + if len(manager.all_conv_states) != len(expected_offsets) or len(manager.all_ssm_states) != len( + expected_offsets + ): + raise ValueError("V2 Mamba state tensors must match the layer-offset table") + first_conv_state = manager.all_conv_states[0] first_ssm_state = manager.all_ssm_states[0] - conv_slot_stride_bytes = _slot_stride_bytes(first_conv_state) - ssm_slot_stride_bytes = _slot_stride_bytes(first_ssm_state) - conv_slot_bytes = int(first_conv_state[0].numel() * first_conv_state.element_size()) - ssm_slot_bytes = int(first_ssm_state[0].numel() * first_ssm_state.element_size()) - num_slots = int(first_ssm_state.shape[0]) - - # V2 coalesces equal-size buffers into slot-major physical pools. The - # SHARED tensor bases include each layer/role's offset within slot 0, while - # stride(0) is the distance to the same buffer in the next physical slot. - # Preserve both pieces: V1's layer-major ``layer * num_slots`` formula does - # not describe this layout. - conv_layer_slot0_addresses = { - int(global_layer_id): int(manager.all_conv_states[offset].data_ptr()) - for global_layer_id, offset in mamba_layer_offsets.items() - } - ssm_layer_slot0_addresses = { - int(global_layer_id): int(manager.all_ssm_states[offset].data_ptr()) - for global_layer_id, offset in mamba_layer_offsets.items() - } + conv_pool = _build_v2_mamba_state_pool(manager.all_conv_states) + ssm_pool = _build_v2_mamba_state_pool(manager.all_ssm_states) + if conv_pool.num_slots != ssm_pool.num_slots: + raise ValueError("V2 Mamba convolution and SSM states must have the same number of slots") d_conv_m1 = manager.conv_state_shape[1] conv_elem_size = first_conv_state.element_size() @@ -197,22 +233,10 @@ def _build_layer_group_for_v2_mamba( return MambaLayerGroup( pool_group_idx=pool_group_idx, mamba_layer_offsets=mamba_layer_offsets, - conv_states=PhysicalPool( - base_address=int(first_conv_state.data_ptr()), - slot_bytes=conv_slot_bytes, - num_slots=num_slots, - slot_stride_bytes=conv_slot_stride_bytes, - ), - ssm_states=PhysicalPool( - base_address=int(first_ssm_state.data_ptr()), - slot_bytes=ssm_slot_bytes, - num_slots=num_slots, - slot_stride_bytes=ssm_slot_stride_bytes, - ), + conv_states=conv_pool, + ssm_states=ssm_pool, conv_section_bytes=conv_section_bytes, ssm_bytes_per_head=ssm_bytes_per_head, - conv_layer_slot0_addresses=conv_layer_slot0_addresses, - ssm_layer_slot0_addresses=ssm_layer_slot0_addresses, ) diff --git a/tensorrt_llm/_torch/disaggregation/resource/page.py b/tensorrt_llm/_torch/disaggregation/resource/page.py index 28537ff9a5bb..ac48bb5c32ec 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/page.py +++ b/tensorrt_llm/_torch/disaggregation/resource/page.py @@ -62,17 +62,22 @@ class PhysicalPool: base_address: int # uint64 slot_bytes: int num_slots: int - # Distance between the starts of adjacent slots. Most pools are densely - # packed, so the stride defaults to the transferable payload size. V2 - # Mamba views point into a coalesced physical slot whose stride can be - # larger than the state payload described by ``slot_bytes``. + # Distance between adjacent slots in this logical pool view. Mamba also + # consumes ``layer_stride_bytes`` to address its per-layer state. The + # defaults preserve V1 Mamba's dense ``[layer][slot][payload]`` layout; + # V2 Mamba overrides both strides for its slot-major coalesced layout. slot_stride_bytes: Optional[int] = None + layer_stride_bytes: Optional[int] = None def __post_init__(self) -> None: if self.slot_stride_bytes is None: self.slot_stride_bytes = self.slot_bytes + if self.layer_stride_bytes is None: + self.layer_stride_bytes = self.num_slots * self.slot_stride_bytes if self.slot_stride_bytes < self.slot_bytes: raise ValueError("slot_stride_bytes must be greater than or equal to slot_bytes") + if self.layer_stride_bytes < self.slot_bytes: + raise ValueError("layer_stride_bytes must be greater than or equal to slot_bytes") def to_dict(self) -> dict: return { @@ -80,6 +85,7 @@ def to_dict(self) -> dict: "slot_bytes": int(self.slot_bytes), "num_slots": int(self.num_slots), "slot_stride_bytes": int(self.slot_stride_bytes), + "layer_stride_bytes": int(self.layer_stride_bytes), } @staticmethod @@ -93,6 +99,11 @@ def from_dict(data: dict) -> "PhysicalPool": if data.get("slot_stride_bytes") is not None else None ), + layer_stride_bytes=( + int(data["layer_stride_bytes"]) + if data.get("layer_stride_bytes") is not None + else None + ), ) @@ -235,11 +246,6 @@ class MambaLayerGroup(LayerGroup): ssm_states: Optional[PhysicalPool] = None conv_section_bytes: Optional[List[int]] = None ssm_bytes_per_head: Optional[int] = None - # V2 pools are slot-major and may coalesce several layer/role buffers into - # one physical slot. These are the manager-provided buffer offsets within - # slot 0; they cannot be derived with V1's layer-major pointer formula. - conv_layer_slot0_addresses: Optional[Dict[int, int]] = None - ssm_layer_slot0_addresses: Optional[Dict[int, int]] = None def to_dict(self) -> dict: return { @@ -249,24 +255,12 @@ def to_dict(self) -> dict: "ssm_states": self.ssm_states.to_dict(), "conv_section_bytes": self.conv_section_bytes, "ssm_bytes_per_head": self.ssm_bytes_per_head, - "conv_layer_slot0_addresses": { - int(k): int(v) for k, v in (self.conv_layer_slot0_addresses or {}).items() - } - if self.conv_layer_slot0_addresses is not None - else None, - "ssm_layer_slot0_addresses": { - int(k): int(v) for k, v in (self.ssm_layer_slot0_addresses or {}).items() - } - if self.ssm_layer_slot0_addresses is not None - else None, } @classmethod def from_dict(cls, data: dict) -> "MambaLayerGroup": conv_section_bytes = data.get("conv_section_bytes") ssm_bytes_per_head = data.get("ssm_bytes_per_head") - conv_layer_slot0_addresses = data.get("conv_layer_slot0_addresses") - ssm_layer_slot0_addresses = data.get("ssm_layer_slot0_addresses") return cls( pool_group_idx=int(data["pool_group_idx"]), mamba_layer_offsets={int(k): int(v) for k, v in data["mamba_layer_offsets"].items()}, @@ -276,14 +270,6 @@ def from_dict(cls, data: dict) -> "MambaLayerGroup": if conv_section_bytes is not None else None, ssm_bytes_per_head=int(ssm_bytes_per_head) if ssm_bytes_per_head is not None else None, - conv_layer_slot0_addresses={ - int(k): int(v) for k, v in conv_layer_slot0_addresses.items() - } - if conv_layer_slot0_addresses is not None - else None, - ssm_layer_slot0_addresses={int(k): int(v) for k, v in ssm_layer_slot0_addresses.items()} - if ssm_layer_slot0_addresses is not None - else None, ) diff --git a/tensorrt_llm/_torch/disaggregation/resource/utils.py b/tensorrt_llm/_torch/disaggregation/resource/utils.py index 874911f81fd9..61ed44bb182c 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/utils.py +++ b/tensorrt_llm/_torch/disaggregation/resource/utils.py @@ -132,11 +132,11 @@ def get_unique_pool_memory_descs( pool_counter = 0 for lg_idx, lg in enumerate(page_table.layer_groups): if isinstance(lg, MambaLayerGroup): - is_v2_layout = ( - lg.conv_layer_slot0_addresses is not None - or lg.ssm_layer_slot0_addresses is not None - ) - if is_v2_layout: + # V2 Mamba layer groups reference manager-owned physical pools. + # V1 Mamba state views are standalone and use the first invalid + # pool-group index after the attention groups. + has_physical_pool_group = 0 <= int(lg.pool_group_idx) < len(page_table.pool_groups) + if has_physical_pool_group: pools_and_sizes = [ (pool, get_pool_bytes(pool)) for pool in page_table.pool_groups[int(lg.pool_group_idx)].pools diff --git a/tests/unittest/disaggregated/region/test_page.py b/tests/unittest/disaggregated/region/test_page.py index 14042299fa3b..5d907381cf76 100644 --- a/tests/unittest/disaggregated/region/test_page.py +++ b/tests/unittest/disaggregated/region/test_page.py @@ -27,15 +27,25 @@ def test_physical_pool_construction(): assert pool.base_address == 0x10000 assert pool.slot_bytes == 256 assert pool.num_slots == 4 + assert pool.slot_stride_bytes == 256 + assert pool.layer_stride_bytes == 1024 def test_physical_pool_roundtrip(): - pool = PhysicalPool(base_address=0x10000, slot_bytes=256, num_slots=4) + pool = PhysicalPool( + base_address=0x10000, + slot_bytes=256, + num_slots=4, + slot_stride_bytes=2048, + layer_stride_bytes=512, + ) d = pool.to_dict() restored = PhysicalPool.from_dict(d) assert restored.base_address == pool.base_address assert restored.slot_bytes == pool.slot_bytes assert restored.num_slots == pool.num_slots + assert restored.slot_stride_bytes == pool.slot_stride_bytes + assert restored.layer_stride_bytes == pool.layer_stride_bytes def test_pool_view_roundtrip(): diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index e31180ce18ec..4549e45f30dd 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -14,10 +14,12 @@ import numpy as np import pytest +import torch from tensorrt_llm._torch.disaggregation.base.region import MemRegionGroup, SpecRegion from tensorrt_llm._torch.disaggregation.resource.kv_extractor import ( KVRegionExtractorV1, + _build_v2_mamba_state_pool, build_page_table, ) from tensorrt_llm._torch.disaggregation.resource.page import MapperKind @@ -255,12 +257,14 @@ def test_mamba_layer_group_serialization(): slot_bytes=128, num_slots=10, slot_stride_bytes=512, + layer_stride_bytes=256, ) ssm_pool = PhysicalPool( base_address=8000, slot_bytes=256, num_slots=8, slot_stride_bytes=1024, + layer_stride_bytes=512, ) mlg = MambaLayerGroup( pool_group_idx=1, @@ -269,17 +273,15 @@ def test_mamba_layer_group_serialization(): ssm_states=ssm_pool, conv_section_bytes=[512, 256, 256], ssm_bytes_per_head=128, - conv_layer_slot0_addresses={10: 1000, 11: 2000, 12: 3000}, - ssm_layer_slot0_addresses={10: 8000, 11: 9000, 12: 10000}, ) d = mlg.to_dict() assert d["mamba_layer_offsets"] == {10: 0, 11: 1, 12: 2} assert d["conv_section_bytes"] == [512, 256, 256] - assert d["conv_layer_slot0_addresses"] == {10: 1000, 11: 2000, 12: 3000} - assert d["ssm_layer_slot0_addresses"] == {10: 8000, 11: 9000, 12: 10000} assert d["conv_states"]["slot_stride_bytes"] == 512 + assert d["conv_states"]["layer_stride_bytes"] == 256 assert d["ssm_states"]["slot_stride_bytes"] == 1024 + assert d["ssm_states"]["layer_stride_bytes"] == 512 from tensorrt_llm._torch.disaggregation.resource.page import LayerGroup @@ -290,18 +292,53 @@ def test_mamba_layer_group_serialization(): assert restored.conv_states.slot_bytes == 128 assert restored.conv_states.num_slots == 10 assert restored.conv_states.slot_stride_bytes == 512 + assert restored.conv_states.layer_stride_bytes == 256 assert get_slot_address(restored.conv_states, 3) == 1000 + 3 * 512 assert restored.ssm_states.base_address == 8000 assert restored.ssm_states.slot_bytes == 256 assert restored.ssm_states.num_slots == 8 assert restored.ssm_states.slot_stride_bytes == 1024 + assert restored.ssm_states.layer_stride_bytes == 512 assert restored.conv_section_bytes == [512, 256, 256] assert restored.ssm_bytes_per_head == 128 - assert restored.conv_layer_slot0_addresses == {10: 1000, 11: 2000, 12: 3000} - assert restored.ssm_layer_slot0_addresses == {10: 8000, 11: 9000, 12: 10000} legacy_pool = PhysicalPool.from_dict({"base_address": 1000, "slot_bytes": 128, "num_slots": 10}) assert legacy_pool.slot_stride_bytes == legacy_pool.slot_bytes + assert legacy_pool.layer_stride_bytes == legacy_pool.num_slots * legacy_pool.slot_bytes + + +def test_v2_mamba_state_pool_uses_affine_layer_and_slot_strides(): + num_slots = 3 + state_bytes = 64 + storage = torch.empty((num_slots, 4, state_bytes), dtype=torch.uint8) + states = [storage[:, 0, :], storage[:, 2, :]] + + pool = _build_v2_mamba_state_pool(states) + + assert pool.base_address == states[0].data_ptr() + assert pool.slot_bytes == state_bytes + assert pool.num_slots == num_slots + assert pool.slot_stride_bytes == 4 * state_bytes + assert pool.layer_stride_bytes == 2 * state_bytes + + +def test_v2_mamba_single_layer_pool_preserves_shared_role_footprint(): + state_bytes = 64 + storage = torch.empty((3, 2, state_bytes), dtype=torch.uint8) + + pool = _build_v2_mamba_state_pool([storage[:, 1, :]]) + + assert pool.slot_bytes == state_bytes + assert pool.slot_stride_bytes == 2 * state_bytes + assert pool.layer_stride_bytes == 2 * state_bytes + + +def test_v2_mamba_state_pool_rejects_non_affine_layer_offsets(): + storage = torch.empty((3, 6, 64), dtype=torch.uint8) + states = [storage[:, 0, :], storage[:, 2, :], storage[:, 5, :]] + + with pytest.raises(ValueError, match="uniform layer stride"): + _build_v2_mamba_state_pool(states) def test_v2_mamba_registration_uses_coalesced_physical_pool(): @@ -331,21 +368,15 @@ def test_v2_mamba_registration_uses_coalesced_physical_pool(): slot_bytes=state_bytes, num_slots=num_slots, slot_stride_bytes=physical_slot_bytes, + layer_stride_bytes=2 * state_bytes, ), ssm_states=PhysicalPool( base_address=1000, slot_bytes=state_bytes, num_slots=num_slots, slot_stride_bytes=physical_slot_bytes, + layer_stride_bytes=2 * state_bytes, ), - conv_layer_slot0_addresses={ - 1: 1000 + state_bytes, - 2: 1000 + state_bytes * 3, - }, - ssm_layer_slot0_addresses={ - 1: 1000, - 2: 1000 + state_bytes * 2, - }, ) page_table = KVCachePageTable( tokens_per_block=16, diff --git a/tests/unittest/disaggregated/test_mamba_transfer.py b/tests/unittest/disaggregated/test_mamba_transfer.py index 0b5c3716c95c..d1967df7cad3 100644 --- a/tests/unittest/disaggregated/test_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_mamba_transfer.py @@ -289,17 +289,17 @@ def test_mamba_policy_slot_major_layer_ptrs(): slot_bytes=10, num_slots=8, slot_stride_bytes=20, + layer_stride_bytes=10, ), ssm_states=page.PhysicalPool( base_address=3000, slot_bytes=20, num_slots=8, slot_stride_bytes=40, + layer_stride_bytes=20, ), conv_section_bytes=[10], ssm_bytes_per_head=10, - conv_layer_slot0_addresses={1: 1000, 2: 1010}, - ssm_layer_slot0_addresses={1: 3000, 2: 3020}, ) peer_mlg = page.MambaLayerGroup( pool_group_idx=0, @@ -309,17 +309,17 @@ def test_mamba_policy_slot_major_layer_ptrs(): slot_bytes=10, num_slots=8, slot_stride_bytes=20, + layer_stride_bytes=10, ), ssm_states=page.PhysicalPool( base_address=7000, slot_bytes=20, num_slots=8, slot_stride_bytes=40, + layer_stride_bytes=20, ), conv_section_bytes=[10], ssm_bytes_per_head=10, - conv_layer_slot0_addresses={1: 5000, 2: 5010}, - ssm_layer_slot0_addresses={1: 7000, 2: 7020}, ) self_ri = rank_info.RankInfo( instance_name="self", @@ -371,6 +371,7 @@ def test_mamba_policy_slot_major_interleaved_role_ptrs(): slot_bytes=state_bytes, num_slots=8, slot_stride_bytes=physical_slot_bytes, + layer_stride_bytes=2 * state_bytes, ) ptrs = peer.MambaPolicy._build_layer_ptrs( @@ -378,10 +379,6 @@ def test_mamba_policy_slot_major_interleaved_role_ptrs(): layer_offsets={1: 0, 2: 1}, overlapping_layers=[1, 2], slot=3, - layer_slot0_addresses={ - 1: 1000 + state_bytes, - 2: 1000 + 3 * state_bytes, - }, ) np.testing.assert_array_equal( From 868e3b75e9f361a9fa06b52ffdf364956e9c9c6d Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 21 Jul 2026 21:35:40 +0800 Subject: [PATCH 06/22] [TRTLLM-11875][refactor] Simplify V2 cache pool roles (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 67 +++++++++---------- .../_torch/pyexecutor/mamba_cache_manager.py | 13 ++-- 2 files changed, 34 insertions(+), 46 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index c3960b4774ea..04e6dbe80df8 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -1114,16 +1114,17 @@ def append_to_kv_heads_per_layer( self._log_kv_cache_pool_lifecycle_mapping() - def _get_pool_page_index_role(self, pool_id: int) -> DataRole: - return Role.KEY + def _get_pool_roles(self, pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: + """Return the roles represented by the two page-table index lanes. - def _get_pool_paired_role(self, pool_id: int) -> Optional[DataRole]: - if self.kv_cache_type == CacheTypeCpp.SELFKONLY: - return None - return Role.VALUE + When present, role B must be addressable from role A using a constant + page-index offset. + """ + role_b = None if self.kv_cache_type == CacheTypeCpp.SELFKONLY else Role.VALUE + return Role.KEY, role_b - def _get_block_scale_role_for_pool(self, pool_id: int) -> Optional[DataRole]: - if self.dtype != DataType.NVFP4 or self._get_pool_page_index_role(pool_id) != Role.KEY: + def _get_block_scale_role(self, role_a: DataRole) -> Optional[DataRole]: + if self.dtype != DataType.NVFP4 or role_a != Role.KEY: return None return Role.KEY_BLOCK_SCALE @@ -1134,17 +1135,17 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: if self.enable_swa_scratch_reuse: for layer_id in typed_range(LayerId(self.num_local_layers)): pool_id = self.impl.get_layer_group_id(layer_id) - page_index_role = self._get_pool_page_index_role(pool_id) + role_a, _ = self._get_pool_roles(pool_id) kv_cache_pool_pointers_list.append( [ self.impl.get_mem_pool_base_address( - layer_id, page_index_role, PageIndexMode.PER_LAYER + layer_id, role_a, PageIndexMode.PER_LAYER ), 0, ] ) if self.dtype == DataType.NVFP4: - block_scale_role = self._get_block_scale_role_for_pool(pool_id) + block_scale_role = self._get_block_scale_role(role_a) block_scale_pool_pointers_list.append( [ self.impl.get_mem_pool_base_address( @@ -1159,17 +1160,15 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: else: for pool_id in range(self.num_pools): layer_id = self.impl.layer_grouping[pool_id][0] - page_index_role = self._get_pool_page_index_role(pool_id) + role_a, _ = self._get_pool_roles(pool_id) kv_cache_pool_pointers_list.append( [ - self.impl.get_mem_pool_base_address( - layer_id, page_index_role, PageIndexMode.SHARED - ), + self.impl.get_mem_pool_base_address(layer_id, role_a, PageIndexMode.SHARED), 0, ] ) if self.dtype == DataType.NVFP4: - block_scale_role = self._get_block_scale_role_for_pool(pool_id) + block_scale_role = self._get_block_scale_role(role_a) block_scale_pool_pointers_list.append( [ self.impl.get_mem_pool_base_address( @@ -1183,23 +1182,21 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: for layer_id in typed_range(LayerId(self.num_local_layers)): layer_group_id = self.impl.get_layer_group_id(layer_id) - page_index_role = self._get_pool_page_index_role(layer_group_id) + role_a, role_b = self._get_pool_roles(layer_group_id) index_base_addr = kv_cache_pool_pointers_list[layer_group_id][0] addr_offset = ( - self.impl.get_mem_pool_base_address( - layer_id, page_index_role, PageIndexMode.SHARED - ) + self.impl.get_mem_pool_base_address(layer_id, role_a, PageIndexMode.SHARED) - index_base_addr ) - offset_divisor = self.impl.get_page_stride(layer_id, page_index_role) - if self._get_pool_paired_role(layer_group_id) is not None: + offset_divisor = self.impl.get_page_stride(layer_id, role_a) + if role_b is not None: offset_divisor *= self.kv_factor offset = exact_div(addr_offset, offset_divisor) if self.dtype != DataType.NVFP4: block_scale_offset = None else: - block_scale_role = self._get_block_scale_role_for_pool(layer_group_id) + block_scale_role = self._get_block_scale_role(role_a) if block_scale_role is None: block_scale_offset = None else: @@ -1211,7 +1208,7 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: - block_scale_base_addr ) block_scale_divisor = self.impl.get_page_stride(layer_id, block_scale_role) - if self._get_pool_paired_role(layer_group_id) is not None: + if role_b is not None: block_scale_divisor *= self.kv_factor block_scale_offset = exact_div(block_scale_addr_offset, block_scale_divisor) @@ -1250,16 +1247,13 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: ) for pool_id in range(self.num_pools): layer_id = self.impl.layer_grouping[pool_id][0] - page_index_role = self._get_pool_page_index_role(pool_id) - self.index_scales[pool_id] = self.impl.get_page_index_scale(layer_id, page_index_role) - paired_role = self._get_pool_paired_role(pool_id) - if paired_role is not None: + role_a, role_b = self._get_pool_roles(pool_id) + self.index_scales[pool_id] = self.impl.get_page_index_scale(layer_id, role_a) + if role_b is not None: self.kv_offset[pool_id] = exact_div( - self.impl.get_mem_pool_base_address(layer_id, paired_role, PageIndexMode.SHARED) - - self.impl.get_mem_pool_base_address( - layer_id, page_index_role, PageIndexMode.SHARED - ), - self.impl.get_page_stride(layer_id, page_index_role), + self.impl.get_mem_pool_base_address(layer_id, role_b, PageIndexMode.SHARED) + - self.impl.get_mem_pool_base_address(layer_id, role_a, PageIndexMode.SHARED), + self.impl.get_page_stride(layer_id, role_a), ) else: self.kv_offset[pool_id] = 0 @@ -1429,11 +1423,10 @@ def _prepare_swa_scratch_copy_tensors(self, index_mapper_capacity: int) -> None: for local_layer_idx in range(self.num_local_layers): layer_id = LayerId(local_layer_idx) pool_id = self.layer_to_pool_mapping_dict[layer_id] - page_index_role = self._get_pool_page_index_role(pool_id) - paired_role = self._get_pool_paired_role(pool_id) + role_a, role_b = self._get_pool_roles(pool_id) roles = [ - page_index_role, - page_index_role if paired_role is None else paired_role, + role_a, + role_a if role_b is None else role_b, ] for role_idx, role in enumerate(roles): converter = self.impl.get_page_index_converter(layer_id, role) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index b3de90d0d779..1da63c1b9c32 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -2690,17 +2690,12 @@ def get_cache_size_per_token(model_config, def _is_local_mamba_layer(self, local_layer_idx: int) -> bool: return self._mamba_layer_mask[self.pp_layers[local_layer_idx]] - def _get_pool_page_index_role(self, pool_id: int) -> DataRole: + def _get_pool_roles(self, + pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: layer_id = int(self.impl.layer_grouping[pool_id][0]) if self._is_local_mamba_layer(layer_id): - return MambaRole.SSM_STATE - return Role.KEY - - def _get_pool_paired_role(self, pool_id: int) -> Optional[DataRole]: - layer_id = int(self.impl.layer_grouping[pool_id][0]) - if self._is_local_mamba_layer(layer_id): - return None - return super()._get_pool_paired_role(pool_id) + return MambaRole.SSM_STATE, None + return super()._get_pool_roles(pool_id) def _max_resident_sequences(self) -> int: return self.max_batch_size * self.mapping.pp_size From ee625df472cc2ef33a666a51157a0d3b4f3858ad Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Wed, 22 Jul 2026 09:23:28 +0800 Subject: [PATCH 07/22] [TRTLLM-11875][fix] Handle Mamba block reuse edge cases (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 9 -- .../_torch/pyexecutor/mamba_cache_manager.py | 27 +++++- .../executor/test_mamba_cache_manager.py | 86 ++++++++++++++++++- 3 files changed, 109 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 23b59ff2f1f5..45e0d2c23ac0 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -146,8 +146,6 @@ def get_kv_cache_manager_cls( "supported with hybrid Mamba / linear-attention models.") state_config = kv_cache_config.mamba_state_config - interval = state_config.periodic_snapshot_interval - has_periodic_snapshots = interval is not None and interval > 0 has_additional_snapshots = bool( state_config.additional_snapshot_offsets_from_start or state_config.additional_snapshot_offsets_from_end) @@ -157,13 +155,6 @@ def get_kv_cache_manager_cls( raise ValueError("Mamba additional snapshot offsets require " "use_kv_cache_manager_v2=True; V1 supports only " "periodic_snapshot_interval.") - if (kv_cache_config.enable_block_reuse and not has_periodic_snapshots - and not has_additional_snapshots): - raise ValueError( - "Hybrid Mamba block reuse requires at least one snapshot " - "policy: set " - "mamba_state_config.periodic_snapshot_interval > 0 or " - "provide additional snapshot offsets.") # Skip Softmax only changes attention kernels. Hybrid models still # need a Mamba-capable cache manager for recurrent state. if is_disagg: diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 1da63c1b9c32..262e0631134c 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -1207,7 +1207,7 @@ def _reset_context_mamba_slots(self, num_contexts: int) -> None: def prepare_expect_snapshot_points(self, requests: List[LlmRequest]) -> None: """Set reusable Mamba snapshot boundaries before scheduling.""" - if not self.kv_cache_config.enable_block_reuse: + if not self.enable_block_reuse: for request in requests: request.expect_snapshot_points = [] return @@ -2328,6 +2328,31 @@ def _setup_state_indices(self, requests=None) -> None: values = self.host_block_offsets[self.recurrent_states_pool_index, rows, 0, bi] invalid_mask = (values < 0) | (values >= max_blocks) + # The C++ recurrent-state manager uses null page-table entries for + # logical blocks that are not snapshot boundaries. Usually a + # context chunk ends exactly at an allocated snapshot (or at the + # final live-state block), but the scheduler may shorten a chunk + # further when KV capacity is tight. In that case the Mamba + # kernel must keep accumulating into the next allocated snapshot + # or final block, matching copyLinearAttentionBlock(), which also + # skips placeholders when it advances the live state. + for bad_i in invalid_mask.nonzero( + as_tuple=False).flatten().tolist(): + req = requests[bad_i] + if req.is_context_finished: + continue + last_prompt_block = (req.prompt_len - + 1) // self.tokens_per_block + row = self.host_block_offsets[self.recurrent_states_pool_index, + rows[bad_i], 0] + candidates = row[block_indices[bad_i]:last_prompt_block + 1] + valid_candidates = ((candidates >= 0) + & (candidates < max_blocks)).nonzero( + as_tuple=False) + if valid_candidates.numel() > 0: + values[bad_i] = candidates[valid_candidates[0, 0]] + + invalid_mask = (values < 0) | (values >= max_blocks) if invalid_mask.any(): bad_i = int(invalid_mask.nonzero(as_tuple=False)[0, 0]) req = requests[bad_i] diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index db6f8aaa1b2d..7257531cce1f 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -319,19 +319,31 @@ def test_hybrid_cache_manager_factory_requires_v2_for_explicit_snapshots( ) -def test_hybrid_cache_manager_factory_requires_snapshot_policy(monkeypatch): +@pytest.mark.parametrize( + ("use_v2", "expected"), + [ + (False, CppMambaHybridCacheManager), + ("auto", CppMambaHybridCacheManager), + (True, V2MambaHybridCacheManager), + ], +) +def test_hybrid_cache_manager_factory_allows_reuse_without_snapshot_policy( + monkeypatch, use_v2, expected +): monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) - with pytest.raises(ValueError, match="at least one snapshot policy"): + assert ( get_kv_cache_manager_cls( _hybrid_model_config(), KvCacheConfig( enable_block_reuse=True, mamba_state_config=MambaStateConfig(periodic_snapshot_interval=0), - use_kv_cache_manager_v2=True, + use_kv_cache_manager_v2=use_v2, ), ) + is expected + ) @pytest.mark.parametrize( @@ -903,6 +915,7 @@ def test_non_mtp_pytorch_prepare_and_get_state_indices_flow(): def test_v2_hybrid_prepare_expect_snapshot_points(): mgr = object.__new__(V2MambaHybridCacheManager) + mgr.enable_block_reuse = True mgr.kv_cache_config = KvCacheConfig( enable_block_reuse=True, mamba_state_config=MambaStateConfig( @@ -928,6 +941,7 @@ def test_v2_hybrid_prepare_expect_snapshot_points(): def test_v2_hybrid_prepare_expect_snapshot_points_without_periodic_snapshots(): mgr = object.__new__(V2MambaHybridCacheManager) + mgr.enable_block_reuse = True mgr.kv_cache_config = KvCacheConfig( enable_block_reuse=True, mamba_state_config=MambaStateConfig( @@ -1312,6 +1326,7 @@ def test_v2_hybrid_pure_mamba_rank_does_not_reserve_attention_page(): def test_cpp_hybrid_prepare_expect_snapshot_points(): mgr = object.__new__(CppMambaHybridCacheManager) + mgr.enable_block_reuse = True mgr.kv_cache_config = KvCacheConfig( enable_block_reuse=True, mamba_state_config=MambaStateConfig(periodic_snapshot_interval=64), @@ -1332,6 +1347,58 @@ def test_cpp_hybrid_prepare_expect_snapshot_points(): ] +@pytest.mark.parametrize( + ("allocated_offsets", "context_current_position", "expected_offset"), + [ + ({9: 25}, 0, 25), + ({7: 17, 9: 25}, 0, 17), + ({7: 17, 9: 25}, 256, 25), + ], +) +def test_cpp_hybrid_state_indices_skip_context_placeholders( + allocated_offsets, context_current_position, expected_offset +): + """Capacity-limited chunks use the next real snapshot/final block.""" + null_index = torch.iinfo(torch.int32).max + block_offsets = torch.full((1, 1, 2, 10), null_index, dtype=torch.int32) + for logical_index, pool_offset in allocated_offsets.items(): + block_offsets[0, 0, 0, logical_index] = pool_offset + + request = SimpleNamespace( + py_request_id=0, + prompt_len=314, + is_context_finished=False, + context_current_position=context_current_position, + context_chunk_size=32, + prepopulated_prompt_len=0, + is_dummy=False, + ) + mgr = object.__new__(CppMambaHybridCacheManager) + mgr.local_num_mamba_layers = 1 + mgr.requests = [request] + mgr.tokens_per_block = 32 + mgr.kv_cache_config = SimpleNamespace(enable_block_reuse=True) + mgr.impl = SimpleNamespace( + copy_batch_block_offsets=lambda *args: None, + get_cache_block_ids=lambda *args: [], + ) + mgr.host_block_offsets = block_offsets + mgr.recurrent_states_pool_index = 0 + mgr.blocks_per_window = {LinearCacheType.RECURRENT_STATES.value: (1264, 0)} + mgr._host_state_indices = torch.zeros(1, dtype=torch.int32) + mgr.cuda_state_indices = torch.zeros(1, dtype=torch.int32) + mgr._row_indices = torch.arange(1, dtype=torch.long) + mgr._request_id_to_state_index = {} + mgr._request_id_to_is_dummy = {} + mgr._dummy_request_mask = None + + mgr._setup_state_indices() + + assert mgr._host_state_indices.tolist() == [expected_offset] + assert mgr.cuda_state_indices.tolist() == [expected_offset] + assert mgr.get_state_indices([request.py_request_id], [False]) == [expected_offset] + + def test_v2_block_reuse_commit_saves_ssm_snapshot_at_snapshot_point(): mgr = object.__new__(V2MambaHybridCacheManager) mgr.enable_block_reuse = True @@ -1386,6 +1453,7 @@ def test_v2_hybrid_add_dummy_requests_forwards_encoder_output_lens(mocker): def test_v2_hybrid_prepare_expect_snapshot_points_clears_when_reuse_disabled(): mgr = object.__new__(V2MambaHybridCacheManager) + mgr.enable_block_reuse = False mgr.kv_cache_config = KvCacheConfig(enable_block_reuse=False) request = SimpleNamespace(prompt_len=150, expect_snapshot_points=[64]) @@ -1396,6 +1464,7 @@ def test_v2_hybrid_prepare_expect_snapshot_points_clears_when_reuse_disabled(): def test_cpp_hybrid_prepare_expect_snapshot_points_clears_when_reuse_disabled(): mgr = object.__new__(CppMambaHybridCacheManager) + mgr.enable_block_reuse = False mgr.kv_cache_config = KvCacheConfig(enable_block_reuse=False) request = SimpleNamespace(prompt_len=150, expect_snapshot_points=[64]) @@ -1404,9 +1473,20 @@ def test_cpp_hybrid_prepare_expect_snapshot_points_clears_when_reuse_disabled(): assert request.expect_snapshot_points == [] +def test_mixed_hybrid_snapshot_hook_clears_when_reuse_disabled(): + mgr = object.__new__(MixedMambaHybridCacheManager) + mgr.enable_block_reuse = False + request = SimpleNamespace(prompt_len=64, expect_snapshot_points=[64]) + + mgr.prepare_expect_snapshot_points([request]) + + assert request.expect_snapshot_points == [] + + @pytest.mark.parametrize("interval", [0, -64, None]) def test_cpp_hybrid_prepare_expect_snapshot_points_clears_for_invalid_interval(interval): mgr = object.__new__(CppMambaHybridCacheManager) + mgr.enable_block_reuse = True mgr.kv_cache_config = SimpleNamespace( enable_block_reuse=True, mamba_state_config=SimpleNamespace(periodic_snapshot_interval=interval), From ae53de644dc34102c4cb751f45583866c528f289 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Wed, 22 Jul 2026 15:39:52 +0800 Subject: [PATCH 08/22] [TRTLLM-11875][fix] Address V2 Mamba review feedback (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- docs/source/features/kvcache.md | 13 +- .../disaggregation/resource/kv_extractor.py | 10 +- .../_torch/disaggregation/transceiver.py | 6 +- tensorrt_llm/_torch/pyexecutor/_util.py | 113 +++++++----- .../_torch/pyexecutor/config_utils.py | 115 ++++++------- .../_torch/pyexecutor/kv_cache_manager_v2.py | 3 +- .../_torch/pyexecutor/kv_cache_transceiver.py | 10 +- .../_torch/pyexecutor/mamba_cache_manager.py | 103 +++++++---- tensorrt_llm/commands/serve.py | 10 +- tensorrt_llm/llmapi/llm_args.py | 95 +++++------ .../kv_cache_manager_v2/_storage_manager.py | 6 +- .../test_disagg_inflight_cancel_gate.py | 6 +- .../executor/test_mamba_cache_manager.py | 161 +++++++++++++----- .../modeling/test_modeling_multimodal.py | 7 +- .../disaggregated/test_mamba_transfer.py | 8 +- .../test_kv_cache_manager_v2.py | 2 +- tests/unittest/llmapi/test_llm_args.py | 45 ++++- 17 files changed, 431 insertions(+), 282 deletions(-) diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index 4768dfd9bc15..8bd239a72106 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -80,13 +80,10 @@ Hybrid Mamba models must retain the recurrent Mamba state together with the attention KV prefix. Snapshot policy is grouped under `kv_cache_config.mamba_state_config`. `periodic_snapshot_interval` controls periodic boundaries. They are disabled by default; set the interval to a -positive value to enable them. The interval is accepted through this nested -configuration in the Python API. YAML/JSON configuration files also accept the -deprecated `kv_cache_config.mamba_state_cache_interval` key and migrate it to -the nested field while loading. Enabling block reuse for a hybrid Mamba model -requires either a positive periodic interval or an explicit V2 snapshot offset; -otherwise, configuration validation rejects the unsupported reuse policy. The -prototype +positive value to enable them. The deprecated +`kv_cache_config.mamba_state_cache_interval` alias remains accepted for +compatibility and is copied to the nested field during validation. New code and +configuration files should use the nested field. The prototype `additional_snapshot_offsets_from_start` and `additional_snapshot_offsets_from_end` options add fixed boundaries. Start offsets count tokens from the beginning of the prompt. End offsets count @@ -106,7 +103,7 @@ kv_cache_config: This retains snapshots after the first 128 tokens, at the end of the prompt, and before the final 32 prompt tokens. Positions outside a particular prompt are ignored. Exact explicit boundaries currently require aggregated serving -with `V2MambaHybridCacheManager`, `max_beam_width=1`, and no KV connector. +with `MambaHybridCacheManagerV2`, `max_beam_width=1`, and no KV connector. Hybrid Mamba models use the V1 C++ compatibility manager by default; select V2 explicitly with `use_kv_cache_manager_v2: true`. V1 and the current disaggregated-serving route support periodic snapshots only, while V2 does not diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index 599fd06bd560..c1a6a28b9c0e 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -38,7 +38,7 @@ from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( MambaHybridCacheManager, - V2MambaHybridCacheManager, + MambaHybridCacheManagerV2, ) from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import get_size_in_bytes, nvtx_range @@ -155,7 +155,7 @@ def _build_layer_group_for_mamba( ) -def _slot_stride_bytes(tensor) -> int: +def _slot_stride_bytes(tensor: torch.Tensor) -> int: return int(tensor.stride(0) * tensor.element_size()) @@ -200,7 +200,7 @@ def _build_v2_mamba_state_pool(states: Sequence[torch.Tensor]) -> PhysicalPool: def _build_layer_group_for_v2_mamba( - manager: V2MambaHybridCacheManager, pool_group_idx: int + manager: MambaHybridCacheManagerV2, pool_group_idx: int ) -> MambaLayerGroup: mamba_layer_offsets = { int(global_layer_id): int(local_layer_id) @@ -447,7 +447,7 @@ def _window_size_for_layer(internal_layer_id: int): for variant in pg_desc.slot_desc.variants: layer_group_id = int(variant.layer_group_id) all_internal_layer_ids = list(manager.impl.layer_grouping[layer_group_id]) - if isinstance(manager, V2MambaHybridCacheManager) and any( + if isinstance(manager, MambaHybridCacheManagerV2) and any( manager._is_local_mamba_layer(int(layer_id)) for layer_id in all_internal_layer_ids ): layer_groups_by_id[layer_group_id] = _build_layer_group_for_v2_mamba( @@ -509,7 +509,7 @@ def _window_size_for_layer(internal_layer_id: int): layer_groups.append(layer_group) if isinstance(manager, MambaHybridCacheManager) and not isinstance( - manager, V2MambaHybridCacheManager + manager, MambaHybridCacheManagerV2 ): mamba_layer_group_idx = len(pool_groups) mamba_layer_group = _build_layer_group_for_mamba(manager, mamba_layer_group_idx) diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 3bcc7ec72d6f..bd2dbdbd44a3 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -32,7 +32,7 @@ from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( MambaHybridCacheManager, - V2MambaHybridCacheManager, + MambaHybridCacheManagerV2, ) from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import nvtx_range @@ -141,7 +141,7 @@ def _exchange_rank_info(self): endpoints = cast(list, self._dist.allgather(self._transfer_worker.sender_endpoint)) layer_num = len(self._kv_cache_manager.pp_layers) if isinstance(self._kv_cache_manager, MambaHybridCacheManager) and not isinstance( - self._kv_cache_manager, V2MambaHybridCacheManager + self._kv_cache_manager, MambaHybridCacheManagerV2 ): layer_num += len(self._kv_cache_manager._impl.mamba_layer_offsets) layer_num_per_pp = cast(list, getattr(self._dist, "pp_allgather")(layer_num)) @@ -235,7 +235,7 @@ def _create_kv_slice(self, req: LlmRequest) -> KVSlice: groups.append(block_ids) mamba_state_index = None - if isinstance(self._kv_cache_manager, V2MambaHybridCacheManager): + if isinstance(self._kv_cache_manager, MambaHybridCacheManagerV2): if self._kv_cache_manager.local_num_mamba_layers > 0: mamba_state_index = self._kv_cache_manager.get_state_indices([req.py_request_id])[0] elif isinstance(self._kv_cache_manager, MambaHybridCacheManager): diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 45e0d2c23ac0..9731e70a3e8c 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -47,9 +47,9 @@ from ..speculative import (get_num_extra_kv_tokens, get_num_spec_layers, get_spec_decoder, should_use_separate_draft_kv_cache) from ..utils import is_gdn_replay_enabled -from .config_utils import (extract_mamba_kv_cache_params, is_gemma4_hybrid, - is_hybrid_linear, is_mla, is_nemotron_hybrid, - is_qwen3_hybrid) +from .config_utils import (MambaKVCacheParams, extract_mamba_kv_cache_params, + is_gemma4_hybrid, is_hybrid_linear, is_mla, + is_nemotron_hybrid, is_qwen3_hybrid) from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager from .guided_decoder import GuidedDecoder @@ -58,8 +58,8 @@ from .llm_request import ExecutorResponse, LlmRequestState from .mamba_cache_manager import (BaseMambaCacheManager, CppMambaHybridCacheManager, + MambaHybridCacheManagerV2, MixedMambaHybridCacheManager, - V2MambaHybridCacheManager, use_py_mamba_cache_manager) from .model_engine import PyTorchModelEngine from .py_executor import PyExecutor @@ -112,7 +112,7 @@ def get_kv_cache_manager_cls( """Resolve the concrete KV cache manager class for ``model_config``. For hybrid mamba models the choice between - ``V2MambaHybridCacheManager`` and compatibility managers is made here. + ``MambaHybridCacheManagerV2`` and compatibility managers is made here. Callers that don't care about disagg can omit ``is_disagg`` and get the unified-pool default. @@ -238,7 +238,7 @@ def get_kv_cache_manager_cls( "V2 Mamba block reuse is not compatible with " "enable_kv_pool_rebalance because the rebalancer does not " "yet model retained recurrent-state snapshots.") - return V2MambaHybridCacheManager + return MambaHybridCacheManagerV2 elif sparse_attn_config is not None: return get_sparse_attn_kv_cache_manager(sparse_attn_config) else: @@ -443,7 +443,7 @@ def _get_model_kv_cache_manager_cls( if is_hybrid_linear(model_engine.model.model_config.pretrained_config) \ and kv_cache_config.enable_block_reuse \ and self._speculative_config is not None: - if not issubclass(cls, V2MambaHybridCacheManager): + if not issubclass(cls, MambaHybridCacheManagerV2): logger.warning( "Block reuse does not work with MTP for hybrid linear models " f"when using non-V2 Mamba cache manager {cls.__name__}") @@ -525,14 +525,6 @@ def _per_manager_cache_cost(self, spec_config=self._speculative_config, **extra_kwargs)) - def _get_separate_target_layer_mask(self) -> Optional[List[bool]]: - """Return the target-only mask used by a separate draft cache.""" - if not self._should_create_separate_draft_kv_cache(): - return None - num_target_layers = (self._model_engine.model.model_config. - pretrained_config.num_hidden_layers) - return [True] * num_target_layers - def _get_one_model_draft_layer_mask(self) -> List[bool]: """Return the same draft-only mask used by runtime construction.""" num_draft_layers = self._get_num_draft_layers() @@ -553,13 +545,13 @@ def _get_kv_size_per_token(self, kv_cache_config = (kv_cache_config if kv_cache_config is not None else self._kv_cache_config) model_config = self._model_engine.model.model_config - target_layer_mask = self._get_separate_target_layer_mask() - target_kwargs = ({ - "layer_mask": target_layer_mask - } if target_layer_mask is not None else {}) - total = self._per_manager_cache_cost(self._kv_cache_manager_cls, - model_config, kv_cache_config, - **target_kwargs) + use_separate_draft_kv_cache = ( + self._should_create_separate_draft_kv_cache()) + total = self._per_manager_cache_cost( + self._kv_cache_manager_cls, + model_config, + kv_cache_config, + use_separate_draft_kv_cache=use_separate_draft_kv_cache) if self._is_encoder_decoder(): total += CacheCost.from_raw(self._get_cross_kv_size_per_token()) if self._draft_model_engine is not None: @@ -569,7 +561,7 @@ def _get_kv_size_per_token(self, total += self._per_manager_cache_cost(draft_kv_cache_manager_cls, draft_model_config, kv_cache_config) - elif self._should_create_separate_draft_kv_cache(): + elif use_separate_draft_kv_cache: # One-model draft with separate KV cache layout. # Pass num_layers explicitly since the HF config may report a # different layer count than what is actually used at runtime @@ -595,7 +587,7 @@ def _get_kv_size_per_token(self, effective_draft_config, kv_cache_config, num_layers=self._get_num_draft_layers(), - layer_mask=self._get_one_model_draft_layer_mask()) + is_draft=True) return total def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction, @@ -1281,13 +1273,13 @@ def _get_target_and_draft_cache_costs( target_kv_cache_config = (kv_cache_config if kv_cache_config is not None else self._kv_cache_config) total_kv = self._get_kv_size_per_token(target_kv_cache_config) - target_layer_mask = self._get_separate_target_layer_mask() - target_kwargs = ({ - "layer_mask": target_layer_mask - } if target_layer_mask is not None else {}) + use_separate_draft_kv_cache = ( + self._should_create_separate_draft_kv_cache()) target_kv = self._per_manager_cache_cost( - self._kv_cache_manager_cls, self._model_engine.model.model_config, - target_kv_cache_config, **target_kwargs) + self._kv_cache_manager_cls, + self._model_engine.model.model_config, + target_kv_cache_config, + use_separate_draft_kv_cache=use_separate_draft_kv_cache) # The draft contribution is whatever the aggregate has on top of the # target. Both pieces are CacheCost; subtraction is component-wise. draft_kv = CacheCost(slope=total_kv.slope - target_kv.slope, @@ -1790,6 +1782,21 @@ def _build_per_layer_num_kv_heads( ] * num_spec_layers +def _get_mamba_cache_layer_masks( + mamba_params: MambaKVCacheParams, + mapping: Mapping, + spec_config: Optional[SpeculativeConfig], + is_draft: bool, +) -> tuple[List[bool], List[bool]]: + use_separate_draft_kv_cache = ( + not mapping.enable_attention_dp + and should_use_separate_draft_kv_cache(spec_config)) + return mamba_params.get_layer_masks( + is_draft=is_draft, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + ) + + def _create_kv_cache_manager( model_engine: Optional[PyTorchModelEngine], kv_cache_manager_cls, @@ -1933,7 +1940,7 @@ def _create_kv_cache_manager( manager_extra_kwargs = {} if issubclass(kv_cache_manager_cls, KVCacheManagerV2): manager_extra_kwargs["enable_stats"] = enable_kv_cache_stats - if issubclass(kv_cache_manager_cls, V2MambaHybridCacheManager): + if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): manager_extra_kwargs["is_disagg"] = is_disagg if is_mla(config): @@ -1975,10 +1982,18 @@ def _create_kv_cache_manager( mamba_params = extract_mamba_kv_cache_params( config, - layer_mask=layer_mask, spec_config=spec_config, quant_config=quant_config, ) + mamba_layer_mask, full_attention_layer_mask = ( + _get_mamba_cache_layer_masks( + mamba_params, + mapping, + spec_config, + is_draft, + )) + num_mamba_layers = (0 if is_draft and mamba_params.num_draft_layers > 0 + else mamba_params.num_mamba_layers) # Replay state update kernel for MTP: default on for sm >= 80; gates # below disable it for incompatible feature combinations. Cpp cache @@ -2038,7 +2053,7 @@ def _create_kv_cache_manager( and mamba_params.mamba_ssm_cache_dtype == torch.float16) mamba_manager_extra_kwargs = dict(manager_extra_kwargs) - if issubclass(kv_cache_manager_cls, V2MambaHybridCacheManager): + if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): mamba_manager_extra_kwargs["conv_state_layout"] = "x_b_c" else: mamba_manager_extra_kwargs["model_type"] = "nemotron_hybrid" @@ -2049,15 +2064,15 @@ def _create_kv_cache_manager( mamba_params.num_heads, mamba_params.n_groups, mamba_params.head_dim, - mamba_params.num_mamba_layers, - mamba_params.mamba_layer_mask, + num_mamba_layers, + mamba_layer_mask, mamba_params.dtype, mamba_params.mamba_ssm_cache_dtype, # kv cache parameters kv_cache_config, tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, - num_layers=mamba_params.num_full_attention_layers, - layer_mask=mamba_params.full_attention_layer_mask, + num_layers=sum(full_attention_layer_mask), + layer_mask=full_attention_layer_mask, num_kv_heads=per_layer_num_kv_heads, head_dim=head_dim, tokens_per_block=tokens_per_block, @@ -2084,10 +2099,18 @@ def _create_kv_cache_manager( ) mamba_params = extract_mamba_kv_cache_params( config, - layer_mask=layer_mask, spec_config=spec_config, quant_config=quant_config, ) + mamba_layer_mask, full_attention_layer_mask = ( + _get_mamba_cache_layer_masks( + mamba_params, + mapping, + spec_config, + is_draft, + )) + num_mamba_layers = (0 if is_draft and mamba_params.num_draft_layers > 0 + else mamba_params.num_mamba_layers) # Replay state update for GDN MTP: mirrors the nemotron_hybrid gating # above, minus the Mamba2-specific stochastic-rounding/Philox gate. # The GDN replay kernel does a plain cast on checkpoint commit, so @@ -2130,16 +2153,16 @@ def _create_kv_cache_manager( # contiguous C++ manager state view. V2 exposes per-layer state views, # so enabling the same path there would fail for partitioned batches. if (use_replay and issubclass(kv_cache_manager_cls, - V2MambaHybridCacheManager)): + MambaHybridCacheManagerV2)): logger.info( - "GDN replay is not supported by V2MambaHybridCacheManager; " + "GDN replay is not supported by MambaHybridCacheManagerV2; " "using the legacy MTP path") use_replay = False logger.info("GDN replay state update: " + ("ENABLED" if use_replay else "DISABLED")) mamba_manager_extra_kwargs = dict(manager_extra_kwargs) - if issubclass(kv_cache_manager_cls, V2MambaHybridCacheManager): + if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): mamba_manager_extra_kwargs["conv_state_layout"] = "q_k_v" else: mamba_manager_extra_kwargs["model_type"] = "qwen3_next" @@ -2150,15 +2173,15 @@ def _create_kv_cache_manager( mamba_params.num_heads, mamba_params.n_groups, mamba_params.head_dim, - mamba_params.num_mamba_layers, - mamba_params.mamba_layer_mask, + num_mamba_layers, + mamba_layer_mask, mamba_params.dtype, mamba_params.mamba_ssm_cache_dtype, # kv cache parameters kv_cache_config, tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, - num_layers=mamba_params.num_full_attention_layers, - layer_mask=mamba_params.full_attention_layer_mask, + num_layers=sum(full_attention_layer_mask), + layer_mask=full_attention_layer_mask, num_kv_heads=per_layer_num_kv_heads, head_dim=head_dim, tokens_per_block=tokens_per_block, diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 5f1e0cbc0a26..efb512899bd1 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import dataclasses from typing import List, Optional @@ -185,18 +188,49 @@ class MambaKVCacheParams: n_groups: int # n_groups | linear_num_key_heads head_dim: int # mamba_head_dim | linear_value_head_dim - # Per-layer masks and counts (trailing entries cover MTP/draft layers, - # which are attention-only and carry no Mamba state). + # Target-layer masks and counts. Appended MTP/draft layers need only a + # count because every draft layer is full attention. mamba_layer_mask: List[bool] - full_attention_layer_mask: List[bool] + target_full_attention_layer_mask: List[bool] num_mamba_layers: int - num_full_attention_layers: int + num_draft_layers: int # Dtypes dtype: torch.dtype # config.torch_dtype mamba_ssm_cache_dtype: Optional[ torch.dtype] # quant_config.mamba_ssm_cache_dtype + def get_layer_masks( + self, + *, + is_draft: bool = False, + use_separate_draft_kv_cache: bool = False, + ) -> tuple[List[bool], List[bool]]: + """Return Mamba and attention masks for one cache manager. + + Target masks use target-model layer indices. Appended one-model draft + layers use indices immediately following the target layers, so a + draft-only manager receives target-sized false prefixes. A combined + manager receives the concatenated target and draft layouts. + """ + target_mamba_mask = list(self.mamba_layer_mask) + target_attention_mask = list(self.target_full_attention_layer_mask) + num_target_layers = len(target_mamba_mask) + num_draft_layers = self.num_draft_layers + draft_attention_mask = [True] * num_draft_layers + + if is_draft and num_draft_layers > 0: + return ( + [False] * (num_target_layers + num_draft_layers), + [False] * num_target_layers + draft_attention_mask, + ) + if use_separate_draft_kv_cache: + return target_mamba_mask, target_attention_mask + return ( + target_mamba_mask + [False] * num_draft_layers, + target_attention_mask + draft_attention_mask, + ) + def get_states_bytes_per_layer(self, mapping) -> int: """Return the total bytes of Mamba state per layer, used for budgeting.""" tp_size = mapping.tp_size if not mapping.enable_attention_dp else 1 @@ -214,48 +248,8 @@ def get_states_bytes_per_layer(self, mapping) -> int: return state_bytes_per_layer -def _nemotron_hybrid_layer_masks(config, layer_mask): - pattern = config.hybrid_override_pattern - if layer_mask is None: - return ([c == "*" for c in pattern], [c == "M" for c in pattern]) - - # One-model speculative decoding: layer_mask may extend past the hybrid - # pattern; treat trailing positions as attention-only draft layers. - full_attn, mamba = [], [] - for i, include in enumerate(layer_mask): - if i < len(pattern): - is_attn = pattern[i] == "*" - is_mamba = pattern[i] == "M" - else: - is_attn, is_mamba = True, False - full_attn.append(is_attn and include) - mamba.append(is_mamba and include) - return full_attn, mamba - - -def _qwen3_hybrid_layer_masks(config, layer_mask): - full_attn, mamba = get_qwen3_hybrid_layer_masks(config) - if layer_mask is None: - return full_attn, mamba - - if len(layer_mask) < len(full_attn): - raise ValueError( - "layer_mask is shorter than the Qwen3 hybrid layer pattern") - base_len = len(full_attn) - new_full_attn, new_mamba = [], [] - for i, include in enumerate(layer_mask): - if i < base_len: - new_full_attn.append(full_attn[i] and include) - new_mamba.append(mamba[i] and include) - else: - new_full_attn.append(include) - new_mamba.append(False) - return new_full_attn, new_mamba - - def extract_mamba_kv_cache_params( config, - layer_mask: Optional[List[bool]] = None, spec_config=None, quant_config=None, ) -> MambaKVCacheParams: @@ -265,13 +259,9 @@ def extract_mamba_kv_cache_params( Args: config: HuggingFace model config of a hybrid Mamba model. - layer_mask: Optional per-layer keep mask used by one-model speculative - decoding. Entries past the underlying hybrid pattern length are - treated as attention-only draft layers. When provided, the caller - is responsible for already including spec layers in the mask. - spec_config: When `layer_mask` is None, used to extend the masks with - MTP/draft attention layers (no Mamba state) so they receive KV - cache entries. + spec_config: Optional speculative-decoding config used to describe + appended attention-only MTP/draft layers separately from target + layers. quant_config: Optional, used only to surface `mamba_ssm_cache_dtype`. Returns: @@ -283,31 +273,26 @@ def extract_mamba_kv_cache_params( num_heads = config.mamba_num_heads n_groups = config.n_groups head_dim = config.mamba_head_dim - full_attn_mask, mamba_mask = _nemotron_hybrid_layer_masks( - config, layer_mask) + pattern = config.hybrid_override_pattern + target_full_attn_mask = [layer_type == "*" for layer_type in pattern] + mamba_mask = [layer_type == "M" for layer_type in pattern] elif is_qwen3_hybrid(config): state_size = config.linear_key_head_dim conv_kernel = config.linear_conv_kernel_dim num_heads = config.linear_num_value_heads n_groups = config.linear_num_key_heads head_dim = config.linear_value_head_dim - full_attn_mask, mamba_mask = _qwen3_hybrid_layer_masks( - config, layer_mask) + target_full_attn_mask, mamba_mask = get_qwen3_hybrid_layer_masks(config) else: raise ValueError( f"{type(config).__name__} is not a supported hybrid Mamba config") - # When no explicit layer_mask is given, extend the masks here so MTP/draft - # layers (attention-only, no Mamba state) get KV cache entries. With an - # explicit layer_mask, the caller already encoded those entries. - if layer_mask is None and spec_config is not None: + num_draft_layers = 0 + if spec_config is not None: # Imported lazily to avoid a circular dependency between # config_utils and tensorrt_llm._torch.speculative. from ..speculative.utils import get_num_spec_layers - num_spec_layers = get_num_spec_layers(spec_config) - if num_spec_layers > 0: - full_attn_mask.extend([True] * num_spec_layers) - mamba_mask.extend([False] * num_spec_layers) + num_draft_layers = get_num_spec_layers(spec_config) or 0 mamba_ssm_cache_dtype = None if quant_config is not None: @@ -325,9 +310,9 @@ def extract_mamba_kv_cache_params( n_groups=n_groups, head_dim=head_dim, mamba_layer_mask=mamba_mask, - full_attention_layer_mask=full_attn_mask, + target_full_attention_layer_mask=target_full_attn_mask, num_mamba_layers=sum(mamba_mask), - num_full_attention_layers=sum(full_attn_mask), + num_draft_layers=num_draft_layers, dtype=resolve_hf_torch_dtype(config) or torch.bfloat16, mamba_ssm_cache_dtype=mamba_ssm_cache_dtype, ) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 04e6dbe80df8..ab6053bb1599 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -1097,8 +1097,7 @@ def append_to_kv_heads_per_layer( # capacity lets the next batch of active requests acquire slots without # waiting for the previous batch's transfers to finish. max_num_sequences = max_batch_size * mapping.pp_size - if num_reserved_index_slots < 0: - raise ValueError("num_reserved_index_slots must be non-negative") + assert num_reserved_index_slots >= 0, "num_reserved_index_slots must be non-negative" index_mapper_capacity = ( max_num_sequences * (2 if is_disagg else 1) + num_reserved_index_slots ) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index 9a2ed0ef52fd..99c75ffec131 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -16,8 +16,8 @@ from .llm_request import LlmRequest from .mamba_cache_manager import (BaseMambaCacheManager, CppMambaHybridCacheManager, - MixedMambaHybridCacheManager, - V2MambaHybridCacheManager) + MambaHybridCacheManagerV2, + MixedMambaHybridCacheManager) from .resource_manager import KVCacheManager CacheTransceiverCpp = tensorrt_llm.bindings.internal.batch_manager.CacheTransceiver @@ -162,16 +162,16 @@ def create_kv_cache_transceiver( # transceiver_runtime == None or "CPP" -> use C++ transceiver (default) # transceiver_runtime == "PYTHON" -> use Python transceiver. # - # V2MambaHybridCacheManager is backed by the Python KVCacheManagerV2 core, + # MambaHybridCacheManagerV2 is backed by the Python KVCacheManagerV2 core, # not the C++ BaseKVCacheManager binding required by CacheTransceiverCpp. is_v2_mamba_hybrid = isinstance(mamba_cache_manager, - V2MambaHybridCacheManager) + MambaHybridCacheManagerV2) use_python_transceiver = ( cache_transceiver_config.transceiver_runtime == "PYTHON") if is_v2_mamba_hybrid and not use_python_transceiver: raise ValueError( - "V2MambaHybridCacheManager requires transceiver_runtime='PYTHON' " + "MambaHybridCacheManagerV2 requires transceiver_runtime='PYTHON' " "with backend='NIXL'; it cannot use the C++ transceiver.") if use_python_transceiver: diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 262e0631134c..16da85210d10 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -1121,6 +1121,8 @@ def _allocate_pool_replay_buffers( self.prev_num_accepted_tokens = None self.cache_buf_idx = None self.mamba_ssm_rand_seed = None + self._dummy_request_mask = None + self._dummy_request_mask_host = None self.old_x = None self.old_B = None self.old_dt = None @@ -1169,6 +1171,20 @@ def _allocate_pool_replay_buffers( ) return True + @torch.inference_mode() + def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: + if self._dummy_request_mask is None: + return + + n = len(is_dummy) + assert n <= self._dummy_request_mask_host.shape[0] + self._dummy_request_mask_host.zero_() + if n > 0: + self._dummy_request_mask_host[:n].copy_( + torch.tensor(is_dummy, dtype=torch.bool)) + self._dummy_request_mask.copy_(self._dummy_request_mask_host, + non_blocking=True) + def _reset_context_mamba_slots(self, num_contexts: int) -> None: if num_contexts == 0: return @@ -1582,27 +1598,32 @@ def _get_local_mamba_cache_layout( mapping: Mapping, *, spec_config=None, - layer_mask: Optional[List[bool]] = None, + is_draft: bool = False, + use_separate_draft_kv_cache: bool = False, ): """Return normalized params and local Mamba/attention layer counts. Cache construction and affine sizing must follow the model's PP layout: partition base layers first, then place appended speculative layers on the - last PP rank. An explicit layer mask is authoritative for separate target - and draft caches. + last PP rank. The normalized params retain target masks and the appended + draft-layer count so estimation selects the same combined or per-manager + layout as runtime. """ from tensorrt_llm._torch.pyexecutor.config_utils import \ extract_mamba_kv_cache_params params = extract_mamba_kv_cache_params( model_config.pretrained_config, - layer_mask=layer_mask, spec_config=spec_config, quant_config=model_config.quant_config, ) + mamba_layer_mask, full_attention_layer_mask = params.get_layer_masks( + is_draft=is_draft, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + ) combined_layer_mask = [ is_mamba or is_attention for is_mamba, is_attention in zip( - params.mamba_layer_mask, params.full_attention_layer_mask) + mamba_layer_mask, full_attention_layer_mask) ] local_layer_indices, _ = get_pp_layers( sum(combined_layer_mask), @@ -1610,9 +1631,9 @@ def _get_local_mamba_cache_layout( spec_config=spec_config, layer_mask=combined_layer_mask, ) - local_mamba_layers = sum(params.mamba_layer_mask[layer_idx] + local_mamba_layers = sum(mamba_layer_mask[layer_idx] for layer_idx in local_layer_indices) - local_attention_layers = sum(params.full_attention_layer_mask[layer_idx] + local_attention_layers = sum(full_attention_layer_mask[layer_idx] for layer_idx in local_layer_indices) return params, local_mamba_layers, local_attention_layers @@ -1629,17 +1650,19 @@ def _estimate_mamba_hybrid_cache_cost( num_reserved_dummy_slots: int, include_explicit_snapshots: bool, cap_partial_attention_snapshots: bool, + is_draft: bool = False, + use_separate_draft_kv_cache: bool = False, **kwargs, ) -> Tuple[int, int]: del num_layers spec_config = kwargs.get("spec_config") - layer_mask = kwargs.get("layer_mask") params, local_mamba_layers, local_attention_layers = ( _get_local_mamba_cache_layout( model_config, mapping, spec_config=spec_config, - layer_mask=layer_mask, + is_draft=is_draft, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, )) attention_slope = (KVCacheManager.get_cache_size_per_token( model_config, @@ -2244,20 +2267,6 @@ def update_mamba_states(self, src_state_indices, num_accepted_draft_tokens, state_indices_d) - @torch.inference_mode() - def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: - if self._dummy_request_mask is None: - return - - n = len(is_dummy) - assert n <= self._dummy_request_mask_host.shape[0] - self._dummy_request_mask_host.zero_() - if n > 0: - self._dummy_request_mask_host[:n].copy_( - torch.tensor(is_dummy, dtype=torch.bool)) - self._dummy_request_mask.copy_(self._dummy_request_mask_host, - non_blocking=True) - def get_num_available_tokens(self, token_num_upper_bound: int, max_num_draft_tokens: int = 0, @@ -2435,8 +2444,6 @@ def _setup_states(self) -> None: def _setup_replay_buffers(self, spec_config) -> None: cache_size = self.all_ssm_states.shape[1] device = self.all_ssm_states.device - self._dummy_request_mask = None - self._dummy_request_mask_host = None if not self._allocate_pool_replay_buffers(spec_config, cache_size, device): return @@ -2451,7 +2458,7 @@ def _setup_replay_buffers(self, spec_config) -> None: ) -class V2MambaHybridCacheManager(KVCacheManagerV2, MambaHybridCacheManager): +class MambaHybridCacheManagerV2(KVCacheManagerV2, MambaHybridCacheManager): """Hybrid Mamba cache manager backed by KVCacheManagerV2. Attention KV pages and Mamba recurrent-state pages are both owned by the @@ -2647,6 +2654,7 @@ def __init__( self.layer_offsets[layer_id] for layer_id in self.mamba_pp_layers ] self._request_id_to_state_index = {} + self._request_id_to_is_dummy = {} state_index_capacity = (self.max_batch_size + self._num_reserved_dummy_slots) @@ -2974,7 +2982,19 @@ def _setup_replay_buffers(self, spec_config) -> None: cache_size = self.all_ssm_states[0].shape[0] assert all(t.shape[0] == cache_size for t in self.all_ssm_states) device = self.all_ssm_states[0].device - self._allocate_pool_replay_buffers(spec_config, cache_size, device) + if not self._allocate_pool_replay_buffers(spec_config, cache_size, + device): + return + + mask_capacity = self._host_state_indices.shape[0] + self._dummy_request_mask = torch.zeros(mask_capacity, + dtype=torch.bool, + device=device) + self._dummy_request_mask_host = torch.zeros( + mask_capacity, + dtype=torch.bool, + pin_memory=prefer_pinned(), + ) def _attention_cache_bytes_per_token(self) -> int: # Mamba layers have zero KV heads, so the generic calculation naturally @@ -3083,6 +3103,7 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): if kv_cache is not None and kv_cache.is_active: self.try_commit_blocks(request, kv_cache) self._request_id_to_state_index.pop(request.py_request_id, None) + self._request_id_to_is_dummy.pop(request.py_request_id, None) super().free_resources(request, pin_on_release) def prepare_resources(self, scheduled_batch: ScheduledRequests): @@ -3119,9 +3140,12 @@ def _setup_state_indices(self, requests: List[LlmRequest]) -> None: self.cuda_state_indices.copy_(self._host_state_indices, non_blocking=True) - for i, req in enumerate(requests): - self._request_id_to_state_index[ - req.py_request_id] = self._host_state_indices[i].item() + is_dummy = [req.is_dummy for req in requests] + self._refresh_dummy_request_mask(is_dummy) + state_values = self._host_state_indices[:n].tolist() + for req, value, dummy in zip(requests, state_values, is_dummy): + self._request_id_to_state_index[req.py_request_id] = value + self._request_id_to_is_dummy[req.py_request_id] = dummy def get_state_indices(self, request_ids: Optional[List[int]] = None, @@ -3134,7 +3158,18 @@ def get_state_indices(self, return [0] * len(request_ids) return self.cuda_state_indices if request_ids is not None: - return [self._request_id_to_state_index[rid] for rid in request_ids] + indices = [ + self._request_id_to_state_index[rid] for rid in request_ids + ] + if is_padding is None: + is_padding = [False] * len(request_ids) + assert len(request_ids) == len(is_padding) + is_dummy = [ + self._request_id_to_is_dummy.get(rid, False) or padding + for rid, padding in zip(request_ids, is_padding) + ] + self._refresh_dummy_request_mask(is_dummy) + return indices return self.cuda_state_indices def get_max_resource_count(self) -> int: @@ -3159,12 +3194,16 @@ def update_mamba_states(self, src_state_indices = self.intermediate_state_indices[:num_gens] if self._use_replay_state_update: + assert self._dummy_request_mask is not None + is_dummy_request = self._dummy_request_mask[ + num_contexts:num_contexts + num_gens] replay_metadata = self.get_replay_state_update_metadata() assert replay_metadata is not None _advance_replay_state( replay_metadata, state_indices_d, num_accepted_tokens[num_contexts:num_contexts + num_gens], + is_dummy_request, ) else: for layer_offset, dst in enumerate(self.all_ssm_states): @@ -3268,6 +3307,8 @@ def shutdown(self): self.prev_num_accepted_tokens = None self.cache_buf_idx = None self.mamba_ssm_rand_seed = None + self._dummy_request_mask = None + self._dummy_request_mask_host = None self.old_x = None self.old_B = None self.old_dt = None diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index d753689c65f4..ae3a823c7a66 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -1190,7 +1190,9 @@ def _serve_llm(): llm_args_extra_dict = {} if extra_llm_api_options is not None: with open(extra_llm_api_options, 'r') as f: - llm_args_extra_dict = yaml.safe_load(f) or {} + llm_args_extra_dict = yaml.safe_load(f) + if not isinstance(llm_args_extra_dict, dict): + raise ValueError("Configuration file root must be a mapping.") extra_allow_request_chat_template = _pop_bool_config_option( llm_args_extra_dict, "allow_request_chat_template") allow_request_chat_template = (allow_request_chat_template @@ -1425,7 +1427,9 @@ def serve_encoder(model: str, host: str, port: int, log_level: str, encoder_args_extra_dict = {} if extra_encoder_options is not None: with open(extra_encoder_options, 'r') as f: - encoder_args_extra_dict = yaml.safe_load(f) or {} + encoder_args_extra_dict = yaml.safe_load(f) + if not isinstance(encoder_args_extra_dict, dict): + raise ValueError("Configuration file root must be a mapping.") extra_allow_request_chat_template = _pop_bool_config_option( encoder_args_extra_dict, "allow_request_chat_template") allow_request_chat_template = (allow_request_chat_template @@ -1543,6 +1547,8 @@ def serve_embedding( if extra_llm_api_options is not None: with open(extra_llm_api_options, 'r') as f: extra_dict = yaml.safe_load(f) + if not isinstance(extra_dict, dict): + raise ValueError("Configuration file root must be a mapping.") llm_args = update_llm_args_with_extra_dict( llm_args, extra_dict, explicit_cli_keys=explicit_cli_keys) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index a8f7ce63a598..c135bcf779b3 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3492,52 +3492,6 @@ class MambaStateConfig(StrictBaseModel): "snapshots require KV cache manager V2.") -def _migrate_legacy_mamba_interval_from_config_file( - config_dict: Dict[str, Any]) -> Dict[str, Any]: - """Migrate the legacy Mamba interval in a YAML/JSON config mapping. - - This intentionally runs only at config-file ingestion boundaries. The - Pydantic models remain strict so direct Python construction accepts only - ``mamba_state_config.periodic_snapshot_interval``. - """ - kv_cache_config = config_dict.get("kv_cache_config") - legacy_field = "mamba_state_cache_interval" - if not isinstance(kv_cache_config, - dict) or legacy_field not in kv_cache_config: - return config_dict - - migrated_config = dict(config_dict) - migrated_kv_cache_config = dict(kv_cache_config) - migrated_config["kv_cache_config"] = migrated_kv_cache_config - - state_config_field = "mamba_state_config" - interval_field = "periodic_snapshot_interval" - if state_config_field in migrated_kv_cache_config: - state_config = migrated_kv_cache_config[state_config_field] - if not isinstance(state_config, dict): - raise ValueError( - "Config file field 'kv_cache_config.mamba_state_config' must " - "be a mapping when using the legacy " - "'kv_cache_config.mamba_state_cache_interval' option.") - if interval_field in state_config: - raise ValueError("Config file cannot set both " - "'kv_cache_config.mamba_state_cache_interval' and " - "'kv_cache_config.mamba_state_config." - "periodic_snapshot_interval'.") - migrated_state_config = dict(state_config) - else: - migrated_state_config = {} - - migrated_state_config[interval_field] = migrated_kv_cache_config.pop( - legacy_field) - migrated_kv_cache_config[state_config_field] = migrated_state_config - logger.warning( - "Config-file option 'kv_cache_config.mamba_state_cache_interval' is " - "deprecated; use 'kv_cache_config.mamba_state_config." - "periodic_snapshot_interval' instead.") - return migrated_config - - @PybindMirror.mirror_pybind_fields(_KvCacheConfig) class KvCacheConfig(StrictBaseModel, PybindMirror): """Configuration for the KV cache.""" @@ -3670,6 +3624,15 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): tokens_per_block: int = Field(default=32, description="The number of tokens per block.") + # This is a pure python field, not a pybind field. It is only for the Pytorch backend. + mamba_state_cache_interval: Optional[NonNegativeInt] = Field( + default=None, + status="deprecated", + telemetry=False, + exclude=True, + description= + "Deprecated alias for mamba_state_config.periodic_snapshot_interval.") + # This is a pure python field, not a pybind field. It is only for the Pytorch backend. mamba_state_config: MambaStateConfig = Field( default_factory=MambaStateConfig, @@ -3824,6 +3787,27 @@ def validate_max_gpu_total_bytes(cls, v: int): "kv_cache_config.max_gpu_total_bytes must be non-negative") return v + @model_validator(mode='after') + def migrate_legacy_mamba_interval(self) -> 'KvCacheConfig': + """Copy the deprecated Mamba interval into its nested replacement.""" + if self.mamba_state_cache_interval is None: + return self + if ("periodic_snapshot_interval" + in self.mamba_state_config.model_fields_set): + raise ValueError("Cannot set both " + "'kv_cache_config.mamba_state_cache_interval' and " + "'kv_cache_config.mamba_state_config." + "periodic_snapshot_interval'.") + logger.warning( + "'kv_cache_config.mamba_state_cache_interval' is deprecated; use " + "'kv_cache_config.mamba_state_config." + "periodic_snapshot_interval' instead.") + self.mamba_state_config = self.mamba_state_config.model_copy( + update={ + "periodic_snapshot_interval": self.mamba_state_cache_interval + }) + return self + @model_validator(mode='after') def validate_disk_cache_config(self): if self.disk_cache_size is not None and self.disk_cache_size > 0: @@ -4415,8 +4399,8 @@ def speculative_model(self) -> Optional[Union[str, Path]]: def from_yaml(cls, yaml_path: Union[str, Path]): with open(yaml_path, "r") as f: config_dict = yaml.safe_load(f) - config_dict = _migrate_legacy_mamba_interval_from_config_file( - config_dict) + if not isinstance(config_dict, dict): + raise ValueError("Configuration file root must be a mapping.") return cls(**config_dict) @field_validator("dtype") @@ -5844,8 +5828,7 @@ def update_llm_args_with_extra_dict( If `explicit_cli_keys` is None, YAML wins on conflicts. """ - llm_args_dict = _migrate_legacy_mamba_interval_from_config_file( - llm_args_dict) + llm_args_dict = dict(llm_args_dict) # CLI scalar -> nested KvCacheConfig field. Callers add the CLI scalar # name to `explicit_cli_keys` to make it win over YAML's same-named @@ -5958,11 +5941,13 @@ def update_llm_args_with_extra_options( if extra_llm_api_options is not None: with open(extra_llm_api_options, 'r') as f: llm_args_dict = yaml.safe_load(f) - llm_args = update_llm_args_with_extra_dict( - llm_args, - llm_args_dict, - extra_llm_api_options, - explicit_cli_keys=explicit_cli_keys) + if not isinstance(llm_args_dict, dict): + raise ValueError("Configuration file root must be a mapping.") + llm_args = update_llm_args_with_extra_dict( + llm_args, + llm_args_dict, + extra_llm_api_options, + explicit_cli_keys=explicit_cli_keys) return llm_args diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py index 6f1f1e9dfffc..22b0b70f55d3 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py @@ -958,7 +958,7 @@ def _compute_slots_for_batch( ssm_lc_idx = self._life_cycles.ssm_life_cycle_id total_ssm_slots = sum(kv.num_ssm_slots for kv in batch.kv_caches) num_request_lineages = len(batch.kv_caches) - # A non-live SSM slot indicates that this batch can retain snapshots. + # An extra SSM slot indicates that this batch can retain snapshots. # Snapshot alignment is unknown while storage is being sized, so once # such capacity exists, reserve one partial attention page per request # lineage. V2 replaces covered partial pages within a lineage, making N @@ -966,7 +966,7 @@ def _compute_slots_for_batch( # the bound merely over-reserves attention pages. The guard keeps the # default num_ssm_slots=1 descriptors used by non-Mamba managers from # paying this Mamba-specific overhead. - has_non_live_ssm_capacity = total_ssm_slots > num_request_lineages + has_extra_ssm_slot = total_ssm_slots > num_request_lineages sys_blocks = batch.system_prompt_length // tokens_per_block for lc_idx, lc in typed_enumerate(self._life_cycles.get()): pg_idx = self.get_pool_group_index(lc_idx) @@ -1011,7 +1011,7 @@ def _compute_slots_for_batch( # Retained partial-page copies are additional physical blocks, not # tokens in KVCacheDesc.capacity, so add their safe upper bound # after token staleness and scratch-sharing calculations. - if has_non_live_ssm_capacity: + if has_extra_ssm_slot: num_slots[pg_idx] += num_request_lineages return num_slots diff --git a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py index f5bdc6de75e4..08a6f4701761 100644 --- a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py +++ b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py @@ -25,7 +25,7 @@ from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( CppMambaHybridCacheManager, - V2MambaHybridCacheManager, + MambaHybridCacheManagerV2, ) from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig @@ -529,7 +529,7 @@ def test_python_nixl_transceiver_accepts_v2_mamba_manager(monkeypatch): constructor = Mock(return_value=expected) fake_module = SimpleNamespace(KvCacheTransceiverV2=constructor) monkeypatch.setitem(sys.modules, "tensorrt_llm._torch.disaggregation.transceiver", fake_module) - manager = object.__new__(V2MambaHybridCacheManager) + manager = object.__new__(MambaHybridCacheManagerV2) result = transceiver_module.create_kv_cache_transceiver( Mock(), Mock(), manager, Mock(), config, manager @@ -542,7 +542,7 @@ def test_python_nixl_transceiver_accepts_v2_mamba_manager(monkeypatch): @pytest.mark.parametrize("runtime", [None, "CPP", "auto"]) def test_cpp_runtime_rejects_v2_mamba_manager(runtime): config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime=runtime) - manager = object.__new__(V2MambaHybridCacheManager) + manager = object.__new__(MambaHybridCacheManagerV2) with pytest.raises(ValueError, match="requires transceiver_runtime='PYTHON'"): transceiver_module.create_kv_cache_transceiver( diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 7257531cce1f..cd195c46d0fe 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -18,6 +18,10 @@ _create_kv_cache_manager, get_kv_cache_manager_cls, ) +from tensorrt_llm._torch.pyexecutor.config_utils import ( + MambaKVCacheParams, + extract_mamba_kv_cache_params, +) from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import BlockReusePolicy, KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import ( @@ -28,11 +32,11 @@ from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( MIN_REPLAY_HISTORY_SIZE, CppMambaHybridCacheManager, + MambaHybridCacheManagerV2, MambaRole, MixedMambaHybridCacheManager, PythonMambaCacheManager, ReplayStateUpdateMetadata, - V2MambaHybridCacheManager, _advance_replay_state, _get_local_mamba_cache_layout, _get_mamba_hybrid_pool_size, @@ -146,11 +150,43 @@ def _hybrid_cache_sizing_model_config(layer_types): return SimpleNamespace(pretrained_config=config, quant_config=None) +def test_mamba_kv_cache_params_separate_target_and_draft_masks(): + model_config = _hybrid_cache_sizing_model_config( + [ + "linear_attention", + "full_attention", + "linear_attention", + "full_attention", + ] + ) + params = extract_mamba_kv_cache_params( + model_config.pretrained_config, + spec_config=MTPDecodingConfig(max_draft_len=1), + ) + + assert params.mamba_layer_mask == [True, False, True, False] + assert params.target_full_attention_layer_mask == [False, True, False, True] + assert params.num_draft_layers == 1 + + assert params.get_layer_masks() == ( + [True, False, True, False, False], + [False, True, False, True, True], + ) + assert params.get_layer_masks(use_separate_draft_kv_cache=True) == ( + [True, False, True, False], + [False, True, False, True], + ) + assert params.get_layer_masks(is_draft=True) == ( + [False, False, False, False, False], + [False, False, False, False, True], + ) + + @pytest.mark.parametrize( ("use_v2", "enable_block_reuse", "expected"), [ - (True, False, V2MambaHybridCacheManager), - (True, True, V2MambaHybridCacheManager), + (True, False, MambaHybridCacheManagerV2), + (True, True, MambaHybridCacheManagerV2), (False, False, CppMambaHybridCacheManager), (False, True, CppMambaHybridCacheManager), ("auto", False, CppMambaHybridCacheManager), @@ -180,7 +216,7 @@ class RecordingCppManager(CppMambaHybridCacheManager): def __init__(self, *args, **kwargs): captured_cpp.update(kwargs) - class RecordingV2Manager(V2MambaHybridCacheManager): + class RecordingV2Manager(MambaHybridCacheManagerV2): def __init__(self, *args, **kwargs): captured_v2.update(kwargs) @@ -196,18 +232,18 @@ def __init__(self, *args, **kwargs): pretrained_config=pretrained_config, quant_config=None, ) - mamba_params = SimpleNamespace( + mamba_params = MambaKVCacheParams( state_size=8, conv_kernel=4, num_heads=4, n_groups=1, head_dim=8, - num_mamba_layers=1, mamba_layer_mask=[True, False], + target_full_attention_layer_mask=[False, True], + num_mamba_layers=1, + num_draft_layers=1, dtype=torch.bfloat16, mamba_ssm_cache_dtype=torch.bfloat16, - num_full_attention_layers=1, - full_attention_layer_mask=[False, True], ) monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "1") monkeypatch.setattr("tensorrt_llm._torch.pyexecutor._util.get_sm_version", lambda: 90) @@ -324,7 +360,7 @@ def test_hybrid_cache_manager_factory_requires_v2_for_explicit_snapshots( [ (False, CppMambaHybridCacheManager), ("auto", CppMambaHybridCacheManager), - (True, V2MambaHybridCacheManager), + (True, MambaHybridCacheManagerV2), ], ) def test_hybrid_cache_manager_factory_allows_reuse_without_snapshot_policy( @@ -373,7 +409,7 @@ def test_hybrid_cache_manager_factory_routes_explicit_snapshots_to_v2( use_kv_cache_manager_v2=True, ), ) - is V2MambaHybridCacheManager + is MambaHybridCacheManagerV2 ) @@ -400,7 +436,7 @@ def test_hybrid_cache_manager_factory_routes_explicit_v2_disagg(monkeypatch, bac backend=backend, transceiver_runtime="PYTHON" ), ) - is V2MambaHybridCacheManager + is MambaHybridCacheManagerV2 ) @@ -551,7 +587,7 @@ def test_hybrid_models_resolve_auto_to_python_transceiver(monkeypatch): def test_v2_disagg_slice_skips_state_index_on_mamba_free_pp_rank(): - manager = object.__new__(V2MambaHybridCacheManager) + manager = object.__new__(MambaHybridCacheManagerV2) manager.local_num_mamba_layers = 0 transceiver = object.__new__(KvCacheTransceiverV2) transceiver._kv_cache_manager = manager @@ -594,7 +630,7 @@ def test_v2_hybrid_incompatibility_fails_without_cpp_fallback( with pytest.raises(NotImplementedError, match=expected): creator._fallback_if_unsupported_kv_cache_manager_v2( - V2MambaHybridCacheManager, model_config, KvCacheConfig() + MambaHybridCacheManagerV2, model_config, KvCacheConfig() ) @@ -914,7 +950,7 @@ def test_non_mtp_pytorch_prepare_and_get_state_indices_flow(): def test_v2_hybrid_prepare_expect_snapshot_points(): - mgr = object.__new__(V2MambaHybridCacheManager) + mgr = object.__new__(MambaHybridCacheManagerV2) mgr.enable_block_reuse = True mgr.kv_cache_config = KvCacheConfig( enable_block_reuse=True, @@ -940,7 +976,7 @@ def test_v2_hybrid_prepare_expect_snapshot_points(): def test_v2_hybrid_prepare_expect_snapshot_points_without_periodic_snapshots(): - mgr = object.__new__(V2MambaHybridCacheManager) + mgr = object.__new__(MambaHybridCacheManagerV2) mgr.enable_block_reuse = True mgr.kv_cache_config = KvCacheConfig( enable_block_reuse=True, @@ -971,7 +1007,7 @@ def test_mamba_snapshot_rule_count_deduplicates_and_filters_unreachable_points() def test_v2_hybrid_snapshot_sizing_scales_with_pp_and_explicit_rules(): - mgr = object.__new__(V2MambaHybridCacheManager) + mgr = object.__new__(MambaHybridCacheManagerV2) mgr.max_batch_size = 4 mgr.mapping = Mapping(world_size=2, rank=0, tp_size=1, pp_size=2) mgr.max_seq_len = 128 @@ -1064,7 +1100,7 @@ def test_hybrid_mtp_layout_honors_explicit_base_partition(): for manager_cls in ( CppMambaHybridCacheManager, - V2MambaHybridCacheManager, + MambaHybridCacheManagerV2, ): cache_cost = manager_cls.get_cache_size_per_token( model_config, @@ -1074,7 +1110,7 @@ def test_hybrid_mtp_layout_honors_explicit_base_partition(): spec_config=spec_config, ) extra_attention_bound = ( - 4096 * (rank + 1) if manager_cls is V2MambaHybridCacheManager else 0 + 4096 * (rank + 1) if manager_cls is MambaHybridCacheManagerV2 else 0 ) assert cache_cost == (64 * (rank + 1), 2400 + extra_attention_bound) @@ -1089,19 +1125,17 @@ def test_hybrid_separate_mtp_draft_estimator_has_no_mamba_state(): ] ) spec_config = MTPDecodingConfig(max_draft_len=1) - target_layer_mask = [True] * 4 - draft_layer_mask = [False] * 4 + [True] for rank in range(2): mapping = Mapping(world_size=2, rank=rank, tp_size=1, pp_size=2) - target_cost = V2MambaHybridCacheManager.get_cache_size_per_token( + target_cost = MambaHybridCacheManagerV2.get_cache_size_per_token( model_config, mapping, max_batch_size=1, kv_cache_config=KvCacheConfig(enable_block_reuse=False), spec_config=spec_config, - layer_mask=target_layer_mask, + use_separate_draft_kv_cache=True, ) assert target_cost == (64, 6496) @@ -1109,19 +1143,19 @@ def test_hybrid_separate_mtp_draft_estimator_has_no_mamba_state(): model_config, mapping, spec_config=spec_config, - layer_mask=draft_layer_mask, + is_draft=True, ) assert local_mamba_layers == 0 assert local_attention_layers == rank - draft_cost = V2MambaHybridCacheManager.get_cache_size_per_token( + draft_cost = MambaHybridCacheManagerV2.get_cache_size_per_token( model_config, mapping, max_batch_size=1, kv_cache_config=KvCacheConfig(enable_block_reuse=False), num_layers=1, spec_config=spec_config, - layer_mask=draft_layer_mask, + is_draft=True, ) assert draft_cost == (64 * rank, 4096 * rank) @@ -1164,7 +1198,7 @@ def test_v2_hybrid_estimator_accounts_for_ssm_slot_attention_bound( enable_attention_dp=enable_attention_dp, ) - assert V2MambaHybridCacheManager.get_cache_size_per_token( + assert MambaHybridCacheManagerV2.get_cache_size_per_token( object(), mapping, max_batch_size=4, @@ -1191,7 +1225,7 @@ def test_v2_hybrid_estimator_reserves_attention_for_each_lineage(monkeypatch): # The one CUDA-graph padding slot is less than the four resident request # lineages, but the conservative attention bound still reserves one page # for every lineage. - assert V2MambaHybridCacheManager.get_cache_size_per_token( + assert MambaHybridCacheManagerV2.get_cache_size_per_token( object(), Mapping(world_size=1, tp_size=1, pp_size=1), max_batch_size=4, @@ -1204,7 +1238,7 @@ def test_v2_hybrid_attention_bound_is_snapshot_alignment_agnostic(): mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) def estimate(interval): - return V2MambaHybridCacheManager.get_cache_size_per_token( + return MambaHybridCacheManagerV2.get_cache_size_per_token( model_config, mapping, max_batch_size=2, @@ -1223,7 +1257,7 @@ def estimate(interval): def test_v2_hybrid_planned_capacity_is_bounded_by_resident_sequences(): - mgr = object.__new__(V2MambaHybridCacheManager) + mgr = object.__new__(MambaHybridCacheManagerV2) mgr.max_batch_size = 4 mgr.mapping = Mapping(world_size=2, rank=0, tp_size=1, pp_size=2) mgr.max_seq_len = 128 @@ -1235,7 +1269,7 @@ def test_v2_hybrid_planned_capacity_is_bounded_by_resident_sequences(): def test_v2_hybrid_typical_batch_uses_num_ssm_slots(): - mgr = object.__new__(V2MambaHybridCacheManager) + mgr = object.__new__(MambaHybridCacheManagerV2) mgr.kv_cache_type = CacheTypeCpp.SELF mgr.head_dim_per_layer = [64, 64] mgr.pp_layers = [0, 1] @@ -1278,7 +1312,7 @@ def test_v2_hybrid_typical_batch_uses_num_ssm_slots(): def test_v2_hybrid_rejects_quota_below_live_state_floor(): - mgr = object.__new__(V2MambaHybridCacheManager) + mgr = object.__new__(MambaHybridCacheManagerV2) mgr.max_batch_size = 2 mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) mgr.local_num_mamba_layers = 1 @@ -1305,7 +1339,7 @@ def test_v2_hybrid_rejects_quota_below_live_state_floor(): def test_v2_hybrid_pure_mamba_rank_does_not_reserve_attention_page(): - mgr = object.__new__(V2MambaHybridCacheManager) + mgr = object.__new__(MambaHybridCacheManagerV2) mgr.max_batch_size = 2 mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) mgr.local_num_mamba_layers = 1 @@ -1400,7 +1434,7 @@ def test_cpp_hybrid_state_indices_skip_context_placeholders( def test_v2_block_reuse_commit_saves_ssm_snapshot_at_snapshot_point(): - mgr = object.__new__(V2MambaHybridCacheManager) + mgr = object.__new__(MambaHybridCacheManagerV2) mgr.enable_block_reuse = True mgr.is_draft = False mgr._augment_tokens_for_block_reuse = lambda tokens, request, start, end: tokens[start:end] @@ -1441,7 +1475,7 @@ def test_v2_block_reuse_commit_saves_ssm_snapshot_at_snapshot_point(): def test_v2_hybrid_add_dummy_requests_forwards_encoder_output_lens(mocker): - mgr = object.__new__(V2MambaHybridCacheManager) + mgr = object.__new__(MambaHybridCacheManagerV2) base_add_dummy_requests = mocker.patch.object( KVCacheManagerV2, "add_dummy_requests", return_value=[] ) @@ -1452,7 +1486,7 @@ def test_v2_hybrid_add_dummy_requests_forwards_encoder_output_lens(mocker): def test_v2_hybrid_prepare_expect_snapshot_points_clears_when_reuse_disabled(): - mgr = object.__new__(V2MambaHybridCacheManager) + mgr = object.__new__(MambaHybridCacheManagerV2) mgr.enable_block_reuse = False mgr.kv_cache_config = KvCacheConfig(enable_block_reuse=False) request = SimpleNamespace(prompt_len=150, expect_snapshot_points=[64]) @@ -1515,7 +1549,7 @@ def test_expect_snapshot_points_binding_round_trip(): @skip_no_cuda def test_v2_hybrid_pool_ratio_controls_allocated_memory(): def allocated_memory(pool_ratio): - mgr = object.__new__(V2MambaHybridCacheManager) + mgr = object.__new__(MambaHybridCacheManagerV2) mgr.kv_cache_type = CacheTypeCpp.SELF mgr.head_dim_per_layer = [64, 64] mgr.pp_layers = [0, 1] @@ -1647,7 +1681,7 @@ def _build_v2_hybrid_with_mamba_layer( dtype=DataType.HALF, conv_state_layout="x_b_c", ): - """Construct a real V2MambaHybridCacheManager.""" + """Construct a real MambaHybridCacheManagerV2.""" mamba_mask = [True] * num_mamba_layers + [False] * num_attention_layers attn_mask = [False] * num_mamba_layers + [True] * num_attention_layers if mapping is None: @@ -1665,7 +1699,7 @@ def _build_v2_hybrid_with_mamba_layer( enable_swa_scratch_reuse=enable_swa_scratch_reuse, dtype="nvfp4" if dtype == DataType.NVFP4 else "auto", ) - return V2MambaHybridCacheManager( + return MambaHybridCacheManagerV2( mamba_d_state=8, mamba_d_conv=4, mamba_num_heads=4, @@ -1751,6 +1785,7 @@ def test_v2_hybrid_allocates_mamba_state_and_dummy_indices(): indices = mgr.get_state_indices([123], [False]) assert len(indices) == 1 assert indices[0] >= 0 + assert mgr._request_id_to_is_dummy[123] assert mgr.cuda_state_indices[0].item() == indices[0] assert mgr.get_ssm_states(0).data_ptr() == mgr.all_ssm_states[0].data_ptr() assert mgr.get_conv_states(0).data_ptr() == mgr.all_conv_states[0].data_ptr() @@ -1929,6 +1964,7 @@ def test_v2_hybrid_free_resources_drops_stale_state_index_mapping(): request = mgr.add_dummy_requests([123], token_nums=[8], is_gen=False)[0] request_id = request.py_request_id assert request_id in mgr._request_id_to_state_index + assert request_id in mgr._request_id_to_is_dummy # Move state-index preparation to another request before freeing the # older one, as happens when an asynchronous transfer finishes late. @@ -1936,6 +1972,7 @@ def test_v2_hybrid_free_resources_drops_stale_state_index_mapping(): mgr.free_resources(request) assert request_id not in mgr._request_id_to_state_index + assert request_id not in mgr._request_id_to_is_dummy finally: mgr.shutdown() @@ -2175,6 +2212,54 @@ def test_v2_hybrid_replay_bookkeeping_matches_checkpoint_predicate(monkeypatch): mgr.shutdown() +def test_v2_hybrid_replay_update_skips_dummy_and_padding_rows(monkeypatch): + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.local_num_mamba_layers = 1 + mgr._request_id_to_state_index = { + 100: 0, + 101: 1, + 102: 2, + 103: 3, + } + mgr._request_id_to_is_dummy = { + 100: False, + 101: False, + 102: True, + 103: False, + } + mgr._dummy_request_mask = torch.zeros(4, dtype=torch.bool) + mgr._dummy_request_mask_host = torch.zeros(4, dtype=torch.bool) + mgr._use_replay_state_update = True + mgr.replay_step_width = 5 + mgr.replay_history_size = 16 + mgr.prev_num_accepted_tokens = torch.full((4,), 13, dtype=torch.int32) + mgr.cache_buf_idx = torch.ones(4, dtype=torch.int32) + mgr.intermediate_state_indices = torch.arange(4, dtype=torch.int32) + mgr.all_ssm_states = [] + mgr.all_conv_states = [torch.empty(0)] + mgr.intermediate_conv_states = torch.empty(0) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.mamba_cache_manager._promote_mamba_state_triton", + lambda *args, **kwargs: None, + ) + + request_ids = [100, 101, 102, 103] + state_indices = torch.tensor( + mgr.get_state_indices(request_ids, [False, False, False, True]), + dtype=torch.int32, + ) + assert mgr._dummy_request_mask.tolist() == [False, False, True, True] + + mgr.update_mamba_states( + SimpleNamespace(num_seqs=4, num_contexts=1), + torch.tensor([1, 3, 3, 3], dtype=torch.int32), + state_indices=state_indices, + ) + + assert mgr.prev_num_accepted_tokens.tolist() == [13, 3, 13, 13] + assert mgr.cache_buf_idx.tolist() == [1, 0, 1, 1] + + @skip_no_cuda @pytest.mark.parametrize( "mamba_ssm_cache_dtype", diff --git a/tests/unittest/_torch/modeling/test_modeling_multimodal.py b/tests/unittest/_torch/modeling/test_modeling_multimodal.py index 9d31c6b810c8..7a57ba500eca 100644 --- a/tests/unittest/_torch/modeling/test_modeling_multimodal.py +++ b/tests/unittest/_torch/modeling/test_modeling_multimodal.py @@ -628,6 +628,7 @@ def get_hybrid_kv_cache_manager( } mamba_params = extract_mamba_kv_cache_params(text_config) + mamba_layer_mask, full_attention_layer_mask = mamba_params.get_layer_masks() if mamba_params.dtype not in dtype_map: raise ValueError( f"Unsupported dtype for hybrid cache manager: " @@ -654,15 +655,15 @@ def get_hybrid_kv_cache_manager( mamba_params.n_groups, mamba_params.head_dim, mamba_params.num_mamba_layers, - mamba_params.mamba_layer_mask, + mamba_layer_mask, mamba_params.dtype, mamba_params.mamba_ssm_cache_dtype, # kv cache parameters (positional) kv_cache_config, tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, # kw-only - num_layers=mamba_params.num_full_attention_layers, - layer_mask=mamba_params.full_attention_layer_mask, + num_layers=sum(full_attention_layer_mask), + layer_mask=full_attention_layer_mask, num_kv_heads=text_config.num_key_value_heads, head_dim=head_dim, tokens_per_block=tokens_per_block, diff --git a/tests/unittest/disaggregated/test_mamba_transfer.py b/tests/unittest/disaggregated/test_mamba_transfer.py index d1967df7cad3..c9aee4d5a3d3 100644 --- a/tests/unittest/disaggregated/test_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_mamba_transfer.py @@ -34,8 +34,8 @@ LlmRequestType, ) from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( + MambaHybridCacheManagerV2, MixedMambaHybridCacheManager, - V2MambaHybridCacheManager, ) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import DataType @@ -207,7 +207,7 @@ def _create_managers( mapping = Mapping( world_size=tp, rank=rank, tp_size=tp, pp_size=1, enable_attention_dp=enable_attention_dp ) - manager_cls = V2MambaHybridCacheManager if use_v2 else MixedMambaHybridCacheManager + manager_cls = MambaHybridCacheManagerV2 if use_v2 else MixedMambaHybridCacheManager manager_kwargs = ( { "is_disagg": True, @@ -249,13 +249,13 @@ def _create_managers( def _mamba_layer_ids(manager): - if isinstance(manager, V2MambaHybridCacheManager): + if isinstance(manager, MambaHybridCacheManagerV2): return manager.mamba_layer_offsets return manager._impl.mamba_layer_offsets def _mamba_state_slot(manager, request_id): - if isinstance(manager, V2MambaHybridCacheManager): + if isinstance(manager, MambaHybridCacheManagerV2): return manager.get_state_indices([request_id])[0] return manager.mamba_cache_index[request_id] diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index 168ae0db3fae..a9f5a3b86898 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -2691,7 +2691,7 @@ def test_ssm_slot_attention_bound_reserves_each_lineage(self): slots = manager._storage._compute_slots_for_batch(batch, self.TOKENS_PER_BLOCK, None) self.assertEqual(slots[ssm_pg], 6) - # Any non-live SSM capacity enables the conservative upper bound of one + # Any extra SSM slot enables the conservative upper bound of one # partial attention page for every request lineage. self.assertEqual(slots[attn_pg], 12) manager.shutdown() diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 64a8a6b94aff..58bcca8fd29d 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -151,6 +151,22 @@ def test_from_yaml_migrates_legacy_mamba_interval(self, tmp_path): assert llm_args.kv_cache_config.mamba_state_config.periodic_snapshot_interval == 64 + def test_from_yaml_rejects_empty_file(self, tmp_path): + yaml_path = tmp_path / "empty.yaml" + yaml_path.write_text("", encoding="utf-8") + + with pytest.raises(ValueError, + match="Configuration file root must be a mapping"): + TorchLlmArgs.from_yaml(yaml_path) + + def test_from_yaml_rejects_non_mapping_root(self, tmp_path): + yaml_path = tmp_path / "list.yaml" + yaml_path.write_text("[]", encoding="utf-8") + + with pytest.raises(ValueError, + match="Configuration file root must be a mapping"): + TorchLlmArgs.from_yaml(yaml_path) + def test_llm_args_with_pydantic_options(self): yaml_content = """ max_batch_size: 16 @@ -557,6 +573,8 @@ def get_model_defaults(cls, llm_args): def test_KvCacheConfig_declaration(): + assert KvCacheConfig().mamba_state_cache_interval is None + assert KvCacheConfig().mamba_state_config.periodic_snapshot_interval == 0 assert KvCacheConfig().kv_cache_event_hash_algo == "auto" assert KvCacheConfig().block_reuse_policy == "all_reusable" assert KvCacheConfig().enable_swa_scratch_reuse is False @@ -709,15 +727,24 @@ def test_KvCacheConfig_allows_periodic_snapshots_with_v1(): assert config.mamba_state_config.periodic_snapshot_interval == 64 -def test_KvCacheConfig_rejects_removed_mamba_interval_field(): - with pytest.raises(ValidationError, match="extra_forbidden"): - KvCacheConfig(mamba_state_cache_interval=64) +def test_KvCacheConfig_migrates_deprecated_mamba_interval(monkeypatch): + warnings_seen = [] + monkeypatch.setattr(llm_args_mod.logger, "warning", + lambda message: warnings_seen.append(message)) - with pytest.raises(ValidationError, match="extra_forbidden"): - TorchLlmArgs( - model=llama_model_path, - kv_cache_config={"mamba_state_cache_interval": 64}, - ) + config = KvCacheConfig(mamba_state_cache_interval=64) + + assert config.mamba_state_cache_interval == 64 + assert config.mamba_state_config.periodic_snapshot_interval == 64 + assert any("mamba_state_cache_interval' is deprecated" in message + for message in warnings_seen) + assert "mamba_state_cache_interval" not in config.model_dump() + + llm_args = TorchLlmArgs( + model=llama_model_path, + kv_cache_config={"mamba_state_cache_interval": 32}, + ) + assert llm_args.kv_cache_config.mamba_state_config.periodic_snapshot_interval == 32 def test_config_file_merge_migrates_legacy_mamba_interval_without_mutating_input( @@ -743,7 +770,7 @@ def test_config_file_merge_migrates_legacy_mamba_interval_without_mutating_input def test_config_file_merge_rejects_legacy_and_new_mamba_intervals(): - with pytest.raises(ValueError, match="cannot set both"): + with pytest.raises(ValueError, match="Cannot set both"): update_llm_args_with_extra_dict( {"model": "dummy"}, { From d46b2d48f694c01c820e25f61a4431005f69bcc8 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Wed, 22 Jul 2026 17:30:17 +0800 Subject: [PATCH 09/22] Support per-conversation V2 Mamba reuse (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- docs/source/features/kvcache.md | 4 +- .../_torch/pyexecutor/mamba_cache_manager.py | 10 ++- tensorrt_llm/llmapi/llm_args.py | 16 +++- .../kv_cache_manager_v2/_core/_kv_cache.py | 42 +++++++--- .../executor/test_mamba_cache_manager.py | 77 ++++++++++++++++++ .../test_kv_cache_manager_v2.py | 80 +++++++++++++++++++ tests/unittest/llmapi/test_llm_args.py | 11 +++ 7 files changed, 225 insertions(+), 15 deletions(-) diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index 8bd239a72106..bebd98f33e84 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -88,7 +88,9 @@ configuration files should use the nested field. The prototype `additional_snapshot_offsets_from_end` options add fixed boundaries. Start offsets count tokens from the beginning of the prompt. End offsets count backward from the prompt end, and an end offset of `0` selects the final -prompt boundary. For example: +prompt boundary. The `per_conversation` block reuse policy disables periodic +Mamba snapshots, so configure one or more explicit stable boundaries (usually +an end offset of `0`) when using it with a hybrid Mamba model. For example: ```yaml kv_cache_config: diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 99d8d13c67f9..ef0c0e5498ca 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -2635,8 +2635,12 @@ def __init__( kv_cache_config = kv_cache_config.model_copy(deep=True) if any(mamba_layer_mask) and kv_cache_config.enable_block_reuse: - # SSM reuse is valid only at explicit snapshot boundaries. - kv_cache_config.block_reuse_policy = BlockReusePolicy.PER_REQUEST.value + block_reuse_policy = BlockReusePolicy( + kv_cache_config.block_reuse_policy) + if block_reuse_policy == BlockReusePolicy.ALL_REUSABLE: + # SSM reuse is valid only at explicit snapshot boundaries. + kv_cache_config.block_reuse_policy = ( + BlockReusePolicy.PER_REQUEST.value) self.kv_cache_config = kv_cache_config super().__init__( @@ -3297,6 +3301,8 @@ def update_context_resources(self, if should_commit: self.try_commit_blocks(request, kv_cache) if request.context_remaining_length == 0: + if self.conversation_manager is not None: + self.conversation_manager.save_drop_plan(request, kv_cache) kv_cache.enable_swa_scratch_reuse = False def shutdown(self): diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 410973e0b72c..ca2440799819 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3772,8 +3772,10 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): "'all_reusable' commits reusable blocks after every context chunk; " "'per_request' commits them only after the final context chunk; " "'per_conversation' uses 'per_request' commits and drops the previous " - "turn's committed SWA-window blocks after the current turn's final context " - "chunk. All reusable blocks remain subject to normal cache eviction. " + "turn's committed SWA-window blocks and Mamba stable-boundary state " + "after the current turn's final context chunk. Periodic Mamba state " + "snapshots are disabled with 'per_conversation'. All reusable blocks " + "remain subject to normal cache eviction. " "Requests without conversation params use 'per_request' behavior. When " "'all_reusable' and SWA scratch reuse are both enabled, only non-scratch " "blocks are committed for reuse.") @@ -3858,6 +3860,16 @@ def migrate_legacy_mamba_interval(self) -> 'KvCacheConfig': }) return self + @model_validator(mode='after') + def disable_periodic_mamba_snapshots_for_conversations( + self) -> 'KvCacheConfig': + """Use only explicit stable boundaries for conversation reuse.""" + if (self.block_reuse_policy == "per_conversation" + and self.mamba_state_config.periodic_snapshot_interval != 0): + self.mamba_state_config = self.mamba_state_config.model_copy( + update={"periodic_snapshot_interval": 0}) + return self + @model_validator(mode='after') def validate_disk_cache_config(self): if self.disk_cache_size is not None and self.disk_cache_size > 0: diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py index 5494449c0bd6..9e20abbaeba2 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py @@ -1047,22 +1047,37 @@ def reuse_scope(self) -> ReuseScope: return self._reuse_scope def plan_committed_block_drop(self) -> PlannedDropHandle | None: - """Plan dropping SWA blocks needed only by the next conversation turn. + """Plan dropping pages needed only by the next conversation turn. The plan covers committed pages in each SWA life cycle's current - attention window. Full-attention and attention-sink blocks are excluded - because later turns may still need them. SSM state is not yet supported. - This must be called after stop_committing(). Returns None without - creating a plan if any required SWA page is unavailable. + attention window and the exact SSM snapshot for the committed prefix. + Full-attention and attention-sink blocks are excluded because later + turns may still need them. This must be called after stop_committing(). + Returns None without creating a plan if any required page is unavailable. """ if self._commit_state != self.CommitState.USER_STOP: raise LogicError("plan_committed_block_drop() requires stop_committing()") - end = self._num_committed_blocks + ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id + matched_blocks: list[Block] = [] + if self.num_committed_tokens > 0: + # Locate pages through the radix tree rather than + # SeqBlock.tree_block: the latter is not guaranteed to identify a + # partial snapshot after reuse. Requiring an exact match keeps the + # preceding conversation plan intact if this turn no longer has a + # complete reusable endpoint. All PP ranks use the same lookup so + # attention-only ranks include the final partial SWA block too. + match = self.manager._match_reuse(self.reuse_scope, self._committed_tokens) + if match.num_tokens != self.num_committed_tokens or not match.blocks: + return None + matched_blocks = match.blocks + elif ssm_lc_id is not None: + return None + + end = BlockOrdinal(len(matched_blocks)) pages_to_drop: list[CommittedPage] = [] for lc_idx, lc in self.manager._life_cycles.items(): if isinstance(lc, SsmLifeCycle): - # TODO: Support recording reusable SSM state pages. continue if lc.window_size is None: continue @@ -1071,9 +1086,7 @@ def plan_committed_block_drop(self) -> PlannedDropHandle | None: ) window_start = min(stale_range.end, end) for ordinal in typed_range(window_start, end): - tree_block = self._blocks[ordinal].tree_block - if tree_block is None: - return None + tree_block = matched_blocks[ordinal] page_ref = tree_block.storage[lc_idx] if page_ref is None: return None @@ -1081,6 +1094,15 @@ def plan_committed_block_drop(self) -> PlannedDropHandle | None: if page is None: return None pages_to_drop.append(page) + + if ssm_lc_id is not None: + page_ref = matched_blocks[-1].storage[ssm_lc_id] + if page_ref is None: + return None + page = page_ref() + if not isinstance(page, SsmCommittedPage): + return None + pages_to_drop.append(page) return PlannedDropHandle(pages_to_drop) # Users promise to not commit any more tokens. For cases where we shouldn't reuse generated tokens diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index ef0c61548565..27fa3caa6893 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -1689,6 +1689,9 @@ def _build_v2_hybrid_with_mamba_layer( use_replay_state_update=False, enable_block_reuse=False, enable_partial_reuse=True, + block_reuse_policy="all_reusable", + periodic_snapshot_interval=0, + additional_snapshot_offsets_from_end=None, enable_attention_dp=False, enable_swa_scratch_reuse=False, dtype=DataType.HALF, @@ -1709,7 +1712,12 @@ def _build_v2_hybrid_with_mamba_layer( max_tokens=512, enable_block_reuse=enable_block_reuse, enable_partial_reuse=enable_partial_reuse, + block_reuse_policy=block_reuse_policy, enable_swa_scratch_reuse=enable_swa_scratch_reuse, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=periodic_snapshot_interval, + additional_snapshot_offsets_from_end=list(additional_snapshot_offsets_from_end or []), + ), dtype="nvfp4" if dtype == DataType.NVFP4 else "auto", ) return MambaHybridCacheManagerV2( @@ -2004,6 +2012,75 @@ def test_v2_hybrid_uses_upstream_min_snapshot_policy(): mgr.shutdown() +@skip_no_cuda +def test_v2_hybrid_preserves_per_conversation_and_disables_periodic_snapshots(): + mgr = _build_v2_hybrid_with_mamba_layer( + enable_block_reuse=True, + block_reuse_policy="per_conversation", + periodic_snapshot_interval=64, + additional_snapshot_offsets_from_end=[0], + ) + try: + assert mgr.block_reuse_policy is BlockReusePolicy.PER_CONVERSATION + assert mgr.conversation_manager is not None + assert mgr.kv_cache_config.mamba_state_config.periodic_snapshot_interval == 0 + assert mgr.kv_cache_manager_py_config.commit_min_snapshot + request = SimpleNamespace(prompt_len=150, expect_snapshot_points=[]) + mgr.prepare_expect_snapshot_points([request]) + assert request.expect_snapshot_points == [150] + finally: + mgr.shutdown() + + +def test_v2_hybrid_saves_conversation_plan_only_after_final_context_chunk(): + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.enable_block_reuse = True + mgr.is_draft = False + mgr.block_reuse_policy = BlockReusePolicy.PER_CONVERSATION + events = [] + mgr._augment_tokens_for_block_reuse = lambda tokens, request, start, end: tokens[start:end] + mgr._mark_context_position_as_history = MagicMock() + mgr.conversation_manager = MagicMock() + mgr.conversation_manager.save_drop_plan.side_effect = lambda request, kv_cache: events.append( + "save" + ) + + request = SimpleNamespace( + py_request_id=7, + is_dummy_request=False, + context_current_position=128, + context_remaining_length=0, + expect_snapshot_points=[128], + prompt_len=128, + is_last_context_chunk=True, + get_tokens=lambda beam_idx: list(range(128)), + ) + kv_cache = SimpleNamespace( + is_active=True, + num_committed_tokens=0, + resize=MagicMock(return_value=True), + enable_swa_scratch_reuse=True, + ) + + def commit(tokens): + events.append("commit") + kv_cache.num_committed_tokens += len(tokens) + + kv_cache.commit = MagicMock(side_effect=commit) + kv_cache.stop_committing = MagicMock(side_effect=lambda: events.append("stop")) + mgr.kv_cache_map = {request.py_request_id: kv_cache} + batch = ScheduledRequests() + batch.append_context_request(request) + + mgr.update_context_resources(batch) + + kv_cache.commit.assert_called_once_with(list(range(128))) + kv_cache.stop_committing.assert_called_once_with() + mgr.conversation_manager.save_drop_plan.assert_called_once_with(request, kv_cache) + assert events == ["commit", "stop", "save"] + assert not kv_cache.enable_swa_scratch_reuse + + @skip_no_cuda def test_v2_hybrid_mamba_state_views_use_logical_slots(): mgr = _build_v2_hybrid_with_mamba_layer(max_batch_size=4, num_mamba_layers=2) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index a9f5a3b86898..19e23ffb36a5 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -2140,6 +2140,86 @@ def test_ssm_reuse_keeps_snapshots_from_multiple_commits(self) -> None: kv4.resume(stream) kv4.close() + def test_ssm_planned_drop_targets_latest_snapshot_with_shared_plans(self) -> None: + """Shared plans drop only their conversation endpoint snapshot.""" + cfg = self._make_ssm_config(tokens_per_block=32) + self.manager = KVCacheManager(cfg) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + prompt = [self.next_token() for _ in range(64)] + + kv_cache = self.manager.create_kv_cache() + kv_cache.resume(stream) + kv_cache.capacity = 32 + kv_cache.commit(prompt[:32]) + kv_cache.capacity = 64 + kv_cache.commit(prompt[32:]) + kv_cache.stop_committing() + first_handle = kv_cache.plan_committed_block_drop() + second_handle = kv_cache.plan_committed_block_drop() + self.assertIsNotNone(first_handle) + self.assertIsNotNone(second_handle) + kv_cache.close() + + self.assertEqual(self.manager.probe_reuse(input_tokens=prompt), 64) + assert first_handle is not None + first_handle.drop() + self.assertEqual(self.manager.probe_reuse(input_tokens=prompt), 64) + assert second_handle is not None + second_handle.drop() + self.assertEqual(self.manager.probe_reuse(input_tokens=prompt), 32) + + empty_cache = self.manager.create_kv_cache() + empty_cache.resume(stream) + empty_cache.stop_committing() + self.assertIsNone(empty_cache.plan_committed_block_drop()) + empty_cache.close() + + def test_ssm_planned_drop_includes_partial_swa_window(self) -> None: + """Hybrid plans include SSM and every partial SWA-window page.""" + cfg = self._make_ssm_config( + tokens_per_block=32, + num_attn_layers=1, + num_ssm_layers=1, + window_size=32, + ) + self.manager = KVCacheManager(cfg) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + prompt = [self.next_token() for _ in range(48)] + + kv_cache = self.manager.create_kv_cache() + kv_cache.resume(stream) + kv_cache.capacity = len(prompt) + kv_cache.commit(prompt) + kv_cache.stop_committing() + drop_handle = kv_cache.plan_committed_block_drop() + self.assertIsNotNone(drop_handle) + + match = self.manager._radix_tree.match( + ReuseScope(), prompt, self.manager.enable_partial_match + ) + self.assertEqual(match.num_tokens, len(prompt)) + attn_lc_id = next(iter(self.manager._life_cycles.attention_life_cycles()))[0] + assert self.manager._life_cycles.ssm_life_cycle_id is not None + ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id + planned_pages = [] + for block in match.blocks: + page_ref = block.storage[attn_lc_id] + assert page_ref is not None + planned_pages.append(unwrap_rawref(page_ref)) + ssm_page_ref = match.blocks[-1].storage[ssm_lc_id] + assert ssm_page_ref is not None + planned_pages.append(unwrap_rawref(ssm_page_ref)) + self.assertTrue(all(page.planned_drop_count == 1 for page in planned_pages)) + planned_pages.clear() + del page_ref, ssm_page_ref + + kv_cache.close() + assert drop_handle is not None + drop_handle.drop() + self.assertEqual(self.manager.probe_reuse(input_tokens=prompt), 0) + def test_ssm_same_block_snapshots_support_monotonic_multi_turn_reuse(self) -> None: cfg = self._make_ssm_config(tokens_per_block=32, enable_partial_reuse=True) self.manager = KVCacheManager(cfg) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 75d54bd53611..5944453eacbb 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -700,6 +700,17 @@ def test_KvCacheConfig_declaration(): assert pybind_config.attention_dp_events_gather_period_ms == 10 assert (KvCacheConfig(block_reuse_policy="per_conversation"). block_reuse_policy == "per_conversation") + per_conversation_config = KvCacheConfig( + block_reuse_policy="per_conversation", + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=64, + additional_snapshot_offsets_from_end=[0], + ), + ) + assert (per_conversation_config.mamba_state_config. + periodic_snapshot_interval == 0) + assert (per_conversation_config.mamba_state_config. + additional_snapshot_offsets_from_end == [0]) with pytest.raises(ValidationError): KvCacheConfig(block_reuse_policy="invalid") From cfde749c978ae60a939aa695ae81e5ebd172b752 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Wed, 22 Jul 2026 23:14:09 +0800 Subject: [PATCH 10/22] [TRTLLM-11875][fix] Support dynamic-tree MTP with KV cache V2 (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 11 ++++- .../_torch/speculative/mtp_dynamic_tree.py | 6 ++- .../test_kv_cache_v2_capacity_only.py | 41 +++++++++++++++- .../_torch/speculative/test_eagle3.py | 47 +++++++++++++++++++ 4 files changed, 101 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index a02a70e7c4e3..3361c4e6854f 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -3423,10 +3423,19 @@ def update_resources( # will be resumed by the scheduler on the next iteration. if not kv_cache.is_active: continue + rewind_len = req.py_rewind_len + if self.is_draft: + runtime_draft_len = req.py_rewind_len + req.py_num_accepted_draft_tokens + # Dynamic-tree draft managers reserve K * max_draft_len slots, + # which can exceed the tree's runtime draft width. Reclaim that + # reserve slack together with rejected draft tokens; otherwise + # it accumulates in the draft KV cache after every generation + # step. Target managers do not allocate this reserve slack. + rewind_len += max(self._kv_reserve_draft_tokens - runtime_draft_len, 0) new_capacity = ( None if req.state in (LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT) - else kv_cache.capacity - req.py_rewind_len + else kv_cache.capacity - rewind_len ) history_length = ( None if self.kv_compression_manages_history else req.max_beam_num_tokens - 1 diff --git a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py index 053d23362f98..20791f971f81 100644 --- a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py @@ -624,7 +624,11 @@ def _relocate_kv_eagerly(self, attn_metadata, batch_size): attn_num_heads, self._kv_head_dim_bytes, cache_mgr.max_total_draft_tokens, - cache_mgr.max_attention_window_vec[0], + # Dynamic-tree MTP currently supports full-attention KV layers. + # Hybrid managers may store a recurrent-state sentinel first, + # while V2 represents full attention as None, so neither form is + # suitable for the integer maxKVCacheLen operator argument. + cache_mgr.max_seq_len, pool_pointers, block_offsets, cache_mgr.max_blocks_per_seq, diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py b/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py index df38d7bba8fc..0a05f393346f 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py @@ -16,18 +16,31 @@ CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType -def _manager(*, is_draft: bool, kv_compression_manages_history: bool = False) -> KVCacheManagerV2: +def _manager( + *, + is_draft: bool, + kv_compression_manages_history: bool = False, + kv_reserve_draft_tokens: int = 0, +) -> KVCacheManagerV2: manager = KVCacheManagerV2.__new__(KVCacheManagerV2) manager.is_draft = is_draft manager.kv_compression_manages_history = kv_compression_manages_history + manager._kv_reserve_draft_tokens = kv_reserve_draft_tokens manager.kv_cache_map = {} return manager -def _request(request_id: int, *, rewind: int = 0, complete: bool = False) -> SimpleNamespace: +def _request( + request_id: int, + *, + rewind: int = 0, + accepted_draft_tokens: int = 0, + complete: bool = False, +) -> SimpleNamespace: return SimpleNamespace( py_request_id=request_id, py_rewind_len=rewind, + py_num_accepted_draft_tokens=accepted_draft_tokens, max_beam_num_tokens=201, state=LlmRequestState.GENERATION_COMPLETE if complete @@ -107,6 +120,30 @@ def test_capacity_only_is_scoped_to_target_manager() -> None: target_cache.resize.assert_called_once_with(253, None) +def test_dynamic_tree_draft_reclaims_reserved_capacity() -> None: + manager = _manager(is_draft=True, kv_reserve_draft_tokens=60) + # The runtime tree used 31 draft positions: 26 rejected and 5 accepted. + request = _request(1, rewind=26, accepted_draft_tokens=5) + cache = _cache() + manager.kv_cache_map[request.py_request_id] = cache + + manager.update_resources(SimpleNamespace(generation_requests=[request])) + + # Rewind 26 rejected positions plus the 60 - 31 reserve slack. + cache.resize.assert_called_once_with(201, 200) + + +def test_dynamic_tree_target_does_not_reclaim_unallocated_reserve() -> None: + manager = _manager(is_draft=False, kv_reserve_draft_tokens=60) + request = _request(1, rewind=26, accepted_draft_tokens=5) + cache = _cache() + manager.kv_cache_map[request.py_request_id] = cache + + manager.update_resources(SimpleNamespace(generation_requests=[request])) + + cache.resize.assert_called_once_with(230, 200) + + def test_capacity_only_completion_preserves_history() -> None: manager = _manager(is_draft=False, kv_compression_manages_history=True) request = _request(1, complete=True) diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 33cf3d635661..c74761615101 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -22,6 +22,8 @@ from tensorrt_llm._torch.pyexecutor.py_executor_creator import \ _extend_full_attention_windows_for_spec_decode from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelSpecMetadata +from tensorrt_llm._torch.speculative.mtp_dynamic_tree import \ + MTPEagleDynamicTreeWorker from tensorrt_llm.executor.request import LoRARequest from tensorrt_llm.llmapi import (CudaGraphConfig, Eagle3DecodingConfig, KvCacheConfig, MoeConfig, MTPDecodingConfig) @@ -30,6 +32,51 @@ sys.path.append(os.path.join(os.path.dirname(__file__), '..')) +def test_mtp_dynamic_tree_relocation_uses_full_attention_window( + monkeypatch: pytest.MonkeyPatch) -> None: + worker = object.__new__(MTPEagleDynamicTreeWorker) + worker._kv_head_dim_bytes = 256 + worker._accepted_draft_indices_tensor = torch.tensor([[0, 1], [2, -1]], + dtype=torch.int32) + worker._num_accepted_tokens_buf = torch.tensor([2, 1], dtype=torch.int32) + + attention_pool_pointers = object() + attention_block_offsets = object() + cache_manager = SimpleNamespace( + num_kv_heads_per_layer=[0, 8, 0, 8], + kv_cache_pool_mapping=[[0, 0], [2, 0], [1, 0], [2, 1]], + kv_cache_pool_pointers=[object(), + object(), attention_pool_pointers], + max_attention_window_vec=[None], + max_seq_len=8192, + max_total_draft_tokens=31, + max_blocks_per_seq=256, + tokens_per_block=32, + ) + attention_metadata = SimpleNamespace( + kv_cache_manager=cache_manager, + kv_lens_cuda=torch.tensor([128, 256], dtype=torch.int32), + kv_cache_block_offsets=[object(), + object(), attention_block_offsets], + ) + update_op = MagicMock() + monkeypatch.setattr( + torch.ops.tensorrt_llm, + "update_kv_cache_draft_token_location_2d", + update_op, + ) + + worker._relocate_kv_eagerly(attention_metadata, batch_size=2) + + update_op.assert_called_once() + args = update_op.call_args.args + assert args[4] == 2 + assert args[5] == 8 + assert args[8] == cache_manager.max_seq_len + assert args[9] is attention_pool_pointers + assert args[10] is attention_block_offsets + + def test_eagle3_draft_kv_cache_uses_full_window_when_draft_has_no_swa() -> None: kv_cache_config = KvCacheConfig(max_attention_window=[128, 131072]) draft_pretrained_config = SimpleNamespace(num_hidden_layers=3) From 00bce92cbb90929bc2357a7d5855858604e1ac7e Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Wed, 22 Jul 2026 23:52:27 +0800 Subject: [PATCH 11/22] [Agent fix] Address Mamba snapshot validation feedback (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_torch/disaggregation/resource/page.py | 16 ++- .../_torch/pyexecutor/mamba_cache_manager.py | 13 +- .../_torch/pyexecutor/model_loader.py | 34 ++++- tensorrt_llm/llmapi/llm_args.py | 2 + .../executor/test_mamba_cache_manager.py | 8 +- tests/unittest/llmapi/test_llm_args.py | 120 ++++++++++++++++++ 6 files changed, 183 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/resource/page.py b/tensorrt_llm/_torch/disaggregation/resource/page.py index ac48bb5c32ec..08ed122de058 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/page.py +++ b/tensorrt_llm/_torch/disaggregation/resource/page.py @@ -59,13 +59,21 @@ class MapperKind(IntEnum): @dataclass class PhysicalPool: + """Affine view of a physical pool over logical layers and slots. + + ``slot_bytes`` is the transferable payload for one ``(layer, slot)`` and + ``num_slots`` is the number of logical slots. The payload address is + ``base_address + layer * layer_stride_bytes + slot * slot_stride_bytes``. + The strides describe the physical layout independently of payload size and + slot count. Their defaults describe dense layer-major storage, where + ``slot_stride_bytes == slot_bytes`` and + ``layer_stride_bytes == num_slots * slot_stride_bytes``. V2 Mamba supplies + both explicitly for its slot-major, role-interleaved pools. + """ + base_address: int # uint64 slot_bytes: int num_slots: int - # Distance between adjacent slots in this logical pool view. Mamba also - # consumes ``layer_stride_bytes`` to address its per-layer state. The - # defaults preserve V1 Mamba's dense ``[layer][slot][payload]`` layout; - # V2 Mamba overrides both strides for its slot-major coalesced layout. slot_stride_bytes: Optional[int] = None layer_stride_bytes: Optional[int] = None diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index ef0c0e5498ca..46fe3feea55c 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -74,7 +74,14 @@ def _get_num_cuda_graph_padding_dummy_slots( spec_config: Optional["DecodingBaseConfig"], max_batch_size: int, ) -> int: - """Return the number of persistent CUDA-graph padding dummy IDs.""" + """Return the number of persistent CUDA-graph padding dummy IDs. + + This is computed before ``ModelEngine`` exists and covers draft lengths + reachable at every batch size, including the zero-length acceptance-rate + fallback. ``ModelEngine._compute_dynamic_draft_len_mapping`` is created + later and covers only configured CUDA-graph batch sizes, so it cannot size + this persistent ID set. + """ if spec_config is None: return 1 @@ -2536,7 +2543,9 @@ def __init__( self.replay_step_width: Optional[int] = ( spec_config.tokens_per_gen_step if spec_config is not None and use_replay_state_update else None) - self.replay_history_size: Optional[int] = self.replay_step_width + self.replay_history_size: Optional[int] = ( + max(MIN_REPLAY_HISTORY_SIZE, self.replay_step_width) + if self.replay_step_width is not None else None) self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding self._seed_rank_offset = _mamba_rank_offset(mapping) self._seed_request_counter = 0 diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 688debc80692..f78844f93701 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -42,7 +42,8 @@ MoeLoadBalancer, maybe_create_moe_load_balancer) from ..virtual_memory import RestoreMode from ..virtual_memory import scope as virtual_memory_scope -from .config_utils import resolve_hf_torch_dtype, resolve_ssm_cache_dtype +from .config_utils import (is_hybrid_linear, resolve_hf_torch_dtype, + resolve_ssm_cache_dtype) _KV_CACHE_MAP = { "fp8": QuantAlgo.FP8.value, @@ -52,6 +53,36 @@ _VALID_KV_CACHE_DTYPES = ("fp8", "nvfp4", "auto") +def _validate_and_adjust_mamba_snapshot_config(config: ModelConfig, + llm_args: TorchLlmArgs) -> None: + """Validate snapshot reuse after the model and V2 setting are resolved.""" + if not is_hybrid_linear(config.pretrained_config): + return + + kv_cache_config = llm_args.kv_cache_config + state_config = kv_cache_config.mamba_state_config + has_additional_snapshots = bool( + state_config.additional_snapshot_offsets_from_start + or state_config.additional_snapshot_offsets_from_end) + if (has_additional_snapshots + and kv_cache_config.use_kv_cache_manager_v2 is not True): + raise ValueError( + "Mamba additional snapshot offsets require " + "kv_cache_config.use_kv_cache_manager_v2=True after resolving " + "the model configuration.") + + has_periodic_snapshots = state_config.periodic_snapshot_interval > 0 + if (kv_cache_config.enable_block_reuse and not has_periodic_snapshots + and not has_additional_snapshots): + logger.warning( + "Disabling KV cache block reuse for the hybrid Mamba model " + "because no Mamba state snapshot policy is configured. Set " + "kv_cache_config.mamba_state_config.periodic_snapshot_interval " + "to a positive value or provide additional snapshot offsets to " + "enable block reuse.") + kv_cache_config.enable_block_reuse = False + + def validate_and_set_mamba_ssm_cache_dtype( config: ModelConfig, mamba_ssm_cache_dtype: str, @@ -398,6 +429,7 @@ def load_config_and_apply_defaults( use_kv_cache_manager_v2 = llm_args.kv_cache_config.use_kv_cache_manager_v2 _resolve_kv_cache_manager_v2_auto(llm_args, model_defaults) + _validate_and_adjust_mamba_snapshot_config(config, llm_args) if use_kv_cache_manager_v2 == "auto": logger.info( "Resolved use_kv_cache_manager_v2='auto' to %s for %s", diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 60824c7b017b..42dee4f518e5 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3573,6 +3573,8 @@ class MambaStateConfig(StrictBaseModel): periodic_snapshot_interval: NonNegativeInt = Field( default=0, + status="prototype", + telemetry=True, description= "The number of tokens between periodic snapshots in the Mamba " "prefix cache. Periodic snapshots are disabled by default; set this " diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 27fa3caa6893..b6a2d36135b2 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -369,7 +369,7 @@ def test_hybrid_cache_manager_factory_requires_v2_for_explicit_snapshots( (True, MambaHybridCacheManagerV2), ], ) -def test_hybrid_cache_manager_factory_allows_reuse_without_snapshot_policy( +def test_hybrid_cache_manager_factory_selects_manager_when_snapshot_reuse_disabled( monkeypatch, use_v2, expected ): monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) @@ -379,7 +379,7 @@ def test_hybrid_cache_manager_factory_allows_reuse_without_snapshot_policy( get_kv_cache_manager_cls( _hybrid_model_config(), KvCacheConfig( - enable_block_reuse=True, + enable_block_reuse=False, mamba_state_config=MambaStateConfig(periodic_snapshot_interval=0), use_kv_cache_manager_v2=use_v2, ), @@ -2280,7 +2280,9 @@ def test_v2_hybrid_replay_buffers_size_by_tokens_per_gen_step(): assert mgr.use_replay_state_update is True assert replay_metadata is not None assert replay_metadata.replay_step_width == spec_config.tokens_per_gen_step - assert replay_metadata.replay_history_size == spec_config.tokens_per_gen_step + assert replay_metadata.replay_history_size == max( + MIN_REPLAY_HISTORY_SIZE, spec_config.tokens_per_gen_step + ) layer_cache = mgr.mamba_layer_cache(0) _assert_replay_layer_cache_uses_history_size( layer_cache, replay_metadata.replay_history_size diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 5944453eacbb..c06e50dd62c7 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -3311,6 +3311,126 @@ def get_preferred_transceiver_runtime(cls, pretrained_config=None): return None +class _NoModelDefaults: + + @classmethod + def get_model_defaults(cls, llm_args): + return {} + + +class TestMambaSnapshotConfigResolution: + + @staticmethod + def _load_config(monkeypatch, args, architecture): + from unittest.mock import MagicMock + + from tensorrt_llm._torch.pyexecutor import \ + model_loader as model_loader_mod + + monkeypatch.setattr( + model_loader_mod.AutoModelForCausalLM, + "_resolve_class", + staticmethod(lambda config: _NoModelDefaults), + ) + fake_loader = MagicMock() + fake_config = MagicMock() + fake_config.pretrained_config.architectures = [architecture] + fake_config.pretrained_config.hybrid_override_pattern = None + fake_loader.load_config.return_value = fake_config + return model_loader_mod.ModelLoader.load_config_and_apply_defaults( + "/tmp/dummy_model", args, fake_loader) + + @staticmethod + def _capture_warnings(monkeypatch): + from tensorrt_llm._torch.pyexecutor import \ + model_loader as model_loader_mod + + messages = [] + monkeypatch.setattr( + model_loader_mod.logger, "warning", + lambda message, *args: messages.append(message % args + if args else message)) + return messages + + @pytest.mark.parametrize( + ("kv_cache_config", "expected_reuse", "expected_warning"), + [ + (KvCacheConfig(), False, True), + ( + KvCacheConfig( + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=64), + use_kv_cache_manager_v2=False, + ), + True, + False, + ), + ( + KvCacheConfig( + mamba_state_config=MambaStateConfig( + additional_snapshot_offsets_from_end=[0]), + use_kv_cache_manager_v2=True, + ), + True, + False, + ), + ( + KvCacheConfig( + block_reuse_policy="per_conversation", + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=64), + use_kv_cache_manager_v2=True, + ), + False, + True, + ), + ], + ids=["none", "periodic", "fixed", "per-conversation-no-fixed"], + ) + def test_hybrid_snapshot_policy_controls_block_reuse( + self, + monkeypatch, + kv_cache_config, + expected_reuse, + expected_warning, + ): + args = TorchLlmArgs(model="/tmp/dummy_model", + kv_cache_config=kv_cache_config) + warnings = self._capture_warnings(monkeypatch) + + self._load_config(monkeypatch, args, "Qwen3NextForCausalLM") + + assert args.kv_cache_config.enable_block_reuse is expected_reuse + assert bool(warnings) is expected_warning + if expected_warning: + assert "no Mamba state snapshot policy" in warnings[0] + + def test_hybrid_fixed_snapshot_rejects_auto_resolved_v1(self, monkeypatch): + args = TorchLlmArgs( + model="/tmp/dummy_model", + kv_cache_config=KvCacheConfig( + mamba_state_config=MambaStateConfig( + additional_snapshot_offsets_from_start=[128]), + use_kv_cache_manager_v2="auto", + ), + ) + + with pytest.raises( + ValueError, + match="use_kv_cache_manager_v2=True after resolving"): + self._load_config(monkeypatch, args, "Qwen3NextForCausalLM") + + def test_non_hybrid_without_snapshot_policy_preserves_block_reuse( + self, monkeypatch): + args = TorchLlmArgs(model="/tmp/dummy_model") + warnings = self._capture_warnings(monkeypatch) + + self._load_config(monkeypatch, args, "LlamaForCausalLM") + + assert args.kv_cache_config.enable_block_reuse is True + assert warnings == [] + + class TestTransceiverRuntimeAutoResolution: """Tests for the transceiver_runtime 'auto' selection mechanism.""" From 893c9f2b66a08ce087ed494fb77306af900d6fa8 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Thu, 23 Jul 2026 01:51:08 +0800 Subject: [PATCH 12/22] [Agent fix] Rebuild hybrid dynamic-tree target masks (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_torch/attention_backend/trtllm.py | 5 ++++ .../_torch/speculative/test_eagle3.py | 26 +++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index a3f4e89d8d85..992eec88b7fd 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -1004,6 +1004,11 @@ def update_spec_dec_param( self.use_spec_decoding = self.is_spec_decoding_enabled self.is_spec_dec_tree = is_spec_dec_tree self.is_spec_dec_dynamic_tree = is_spec_dec_dynamic_tree + # A hybrid model's first executed attention layer can have a nonzero + # cache-local index because recurrent layers precede it. Do not rely + # on the C++ ``layer_idx == 0`` fallback to rebuild the target mask: + # the dynamic draft loop clears that mask before the next target step. + self.force_prepare_spec_dec_tree_mask = is_spec_dec_dynamic_tree # Forward static tree length to FMHA kernel selection. self.max_total_draft_tokens = max_total_draft_tokens diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index c74761615101..33a3667d0db4 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -32,6 +32,32 @@ sys.path.append(os.path.join(os.path.dirname(__file__), '..')) +def test_dynamic_tree_metadata_forces_target_mask_prepare_each_step() -> None: + metadata = TrtllmAttentionMetadata( + seq_lens=None, + seq_lens_kv=None, + num_contexts=0, + max_num_requests=1, + max_num_tokens=1, + max_seq_len=1, + ) + common_kwargs = dict( + batch_size=0, + is_spec_decoding_enabled=False, + is_spec_dec_tree=True, + max_draft_len=1, + max_total_draft_tokens=1, + ) + + metadata.update_spec_dec_param(is_spec_dec_dynamic_tree=True, + **common_kwargs) + assert metadata.force_prepare_spec_dec_tree_mask + + metadata.update_spec_dec_param(is_spec_dec_dynamic_tree=False, + **common_kwargs) + assert not metadata.force_prepare_spec_dec_tree_mask + + def test_mtp_dynamic_tree_relocation_uses_full_attention_window( monkeypatch: pytest.MonkeyPatch) -> None: worker = object.__new__(MTPEagleDynamicTreeWorker) From b1a914197001d1031d00eaea238f163d7ee2ef73 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Thu, 23 Jul 2026 01:53:47 +0800 Subject: [PATCH 13/22] [Agent fix] Cover V2 hybrid MTP accuracy (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../defs/accuracy/test_llm_api_pytorch.py | 20 ++++++++++++++++--- .../test_lists/qa/llm_function_core.txt | 6 ++++-- .../test_lists/qa/llm_function_rtx6k.txt | 3 ++- .../test_lists/test-db/l0_b200.yml | 3 ++- .../test_lists/test-db/l0_dgx_b200.yml | 3 ++- tests/integration/test_lists/waives.txt | 6 ++++-- 6 files changed, 31 insertions(+), 10 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index a412513137bf..1689bcf3adde 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -6241,9 +6241,15 @@ def test_bf16(self, moe_backend, tp_size, mocker): extra_acc_spec=extra_acc_spec, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) - def test_bf16_mtp(self, mocker): + @pytest.mark.parametrize( + "v2_kv_cache", + [False, True], + ids=["v1_kv_cache", "v2_kv_cache"], + ) + def test_bf16_mtp(self, v2_kv_cache, mocker): kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.8, - enable_block_reuse=False) + enable_block_reuse=False, + use_kv_cache_manager_v2=v2_kv_cache) cuda_graph_config = CudaGraphConfig(enable_padding=True, max_batch_size=32) @@ -6262,6 +6268,7 @@ def test_bf16_mtp(self, mocker): kv_cache_config=kv_cache_config, cuda_graph_config=cuda_graph_config, speculative_config=mtp_config) as llm: + assert llm.args.kv_cache_config.use_kv_cache_manager_v2 is v2_kv_cache mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN) task = GSM8K(self.MODEL_NAME) @@ -7164,7 +7171,12 @@ def test_nvfp4_parallelism(self, tp_size, ep_size, pp_size, attention_dp): @skip_pre_blackwell @pytest.mark.skip_less_mpi_world_size(8) - def test_nvfp4_8gpus_mtp(self): + @pytest.mark.parametrize( + "v2_kv_cache", + [False, True], + ids=["v1_kv_cache", "v2_kv_cache"], + ) + def test_nvfp4_8gpus_mtp(self, v2_kv_cache): # Test MTP (Multi-Token Prediction) accuracy with nvfp4-fp8kv model. # This test uses MTP with max_draft_len=3 and one_model mode. mtp_config = MTPDecodingConfig( @@ -7178,6 +7190,7 @@ def test_nvfp4_8gpus_mtp(self): enable_block_reuse=True, mamba_ssm_cache_dtype="float16", free_gpu_memory_fraction=0.5, + use_kv_cache_manager_v2=v2_kv_cache, ), max_batch_size=32, tensor_parallel_size=8, @@ -7190,6 +7203,7 @@ def test_nvfp4_8gpus_mtp(self): moe_config=MoeConfig(backend="CUTLASS"), decoding_config=mtp_config, ) as llm: + assert llm.args.kv_cache_config.use_kv_cache_manager_v2 is v2_kv_cache task = MMLU(self.MODEL_NAME) task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index d675ad632e6c..a1e7cac4b001 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -681,7 +681,8 @@ accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_online_e accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_static_eplb[moe_backend=CUTLASS] accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_static_eplb[moe_backend=TRTLLM] accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_static_eplb[moe_backend=CUTEDSL] -accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp +accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp[v1_kv_cache] +accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp[v2_kv_cache] accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp_custom_op accuracy/test_llm_api_pytorch.py::TestNemotronV3Nano::test_nvfp4_marlin_multi_gpus[tp_size=8] accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_marlin_multi_gpus[tp_size=8] @@ -702,7 +703,8 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4[latency_moe_tr accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4[latency_moe_trtllm_attention_dp] accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4_4gpus[latency_moe_cutlass] accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4_4gpus[latency_moe_trtllm_eagle3] -accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp +accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v1_kv_cache] +accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v2_kv_cache] accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_dummy_load_format accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_fp8[latency-torch_compile=False] accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_fp8[latency-torch_compile=True] diff --git a/tests/integration/test_lists/qa/llm_function_rtx6k.txt b/tests/integration/test_lists/qa/llm_function_rtx6k.txt index 4260a5328d77..d16dc73c5745 100644 --- a/tests/integration/test_lists/qa/llm_function_rtx6k.txt +++ b/tests/integration/test_lists/qa/llm_function_rtx6k.txt @@ -173,7 +173,8 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[latency_moe_cutl accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[latency_moe_cutlass-torch_compile=True] accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp1-CUTLASS] accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp1-TRTLLM] -accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp +accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v1_kv_cache] +accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v2_kv_cache] accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8[enable_block_reuse=False] accuracy/test_llm_api_pytorch.py::TestQwen3_5_397B_A17B::test_nvfp4[adp4_cutedsl] accuracy/test_llm_api_pytorch.py::TestQwen3_5_397B_A17B::test_nvfp4[adp4_trtllm] diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index d812dcf5d84e..8a6892bb4491 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -61,7 +61,8 @@ l0_b200: - accuracy/test_llm_api_pytorch_multimodal.py::TestNanoV3Omni::test_auto_dtype[fp8_mmmu_encoder_cuda_graph] - accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[qwen3vl_2b_instruct] - accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[nemotron_nano_v3_omni_nvfp4] - - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp + - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v1_kv_cache] + - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v2_kv_cache] - disaggregated/test_workers.py::test_workers_kv_cache_aware_router_eviction[TinyLlama-1.1B-Chat-v1.0] # nvbugs 5300551 - test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-NVFP4-nvfp4-quantized/Meta-Llama-3.1-8B] - test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-FP8-llama-3.1-model/Llama-3.1-8B-Instruct-FP8] diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 0a8988c17edc..f006349d75d2 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -169,7 +169,8 @@ l0_dgx_b200: - accuracy/test_disaggregated_serving.py::TestKimiK25::test_nvfp4 TIMEOUT (180) - accuracy/test_llm_api_pytorch.py::TestKimiK25::test_nvfp4[tp8] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4ProDSpark::test_gsm8k_dep8_megamoe_deepgemm TIMEOUT (240) - - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp TIMEOUT (60) + - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp[v1_kv_cache] TIMEOUT (60) + - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp[v2_kv_cache] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_8gpus[attention_dp_on-cutedsl] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_8gpus[attention_dp_off-trtllm] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=0-block_reuse=False-use_py_transceiver=True] TIMEOUT (60) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 9cd067100472..45c0e23bb7ed 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -277,7 +277,8 @@ full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_2gpus[ep2-t full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_2gpus[tp2-trtllm-auto] SKIP (https://nvbugs/6273846) full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_2gpus[tp2-trtllm-fp8] SKIP (https://nvbugs/6273846) full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp1-CUTLASS] SKIP (https://nvbugs/6273850) -full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp SKIP (https://nvbugs/6275856) +full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v1_kv_cache] SKIP (https://nvbugs/6275856) +full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v2_kv_cache] SKIP (https://nvbugs/6275856) full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8[enable_block_reuse=False] SKIP (https://nvbugs/6313076) full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_dflash SKIP (https://nvbugs/6273850) full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_fp8 SKIP (https://nvbugs/6273850) @@ -309,7 +310,8 @@ full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::Tes full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_2gpus[tp2-trtllm-auto] SKIP (https://nvbugs/6273846) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_2gpus[tp2-trtllm-fp8] SKIP (https://nvbugs/6273846) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp1-CUTLASS] SKIP (https://nvbugs/6273850) -full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp SKIP (https://nvbugs/6275856) +full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v1_kv_cache] SKIP (https://nvbugs/6275856) +full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v2_kv_cache] SKIP (https://nvbugs/6275856) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8[enable_block_reuse=False] SKIP (https://nvbugs/6313076) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16 SKIP (https://nvbugs/6273850) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_dflash SKIP (https://nvbugs/6273850) From 82b2c8f71f31efab200c7671ce999df67325d5a0 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Thu, 23 Jul 2026 02:24:11 +0800 Subject: [PATCH 14/22] [Agent fix] Handle attention-free Mamba disagg ranks (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../native/mixers/attention/peer.py | 5 +++ .../_torch/disaggregation/native/rank_info.py | 4 +++ tests/unittest/disaggregated/test_peer.py | 31 ++++++++++++++++ .../unittest/disaggregated/test_rank_info.py | 35 +++++++++++++++++++ 4 files changed, 75 insertions(+) diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py index 12d7cadfee4f..dfd18cc19838 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py @@ -395,6 +395,11 @@ def head_match(self, peer_ri: RankInfo) -> tuple[bool, bool]: return head_match, is_dup_head def duplicate_head_factors(self, peer_ri: RankInfo) -> tuple[int, int]: + # Head duplication only applies to attention KV. In particular, do + # not divide 0 / 0 for an attention-free hybrid PP stage; MambaPolicy + # computes its TP mapping separately. + if self._ri.attention.kv_heads_per_rank == 0 or peer_ri.attention.kv_heads_per_rank == 0: + return 1, 1 factor_self, factor_peer = self._head_factors(peer_ri) dup_head = max(1, factor_self // factor_peer) peer_dup_head = max(1, factor_peer // factor_self) diff --git a/tensorrt_llm/_torch/disaggregation/native/rank_info.py b/tensorrt_llm/_torch/disaggregation/native/rank_info.py index 4bea2ee1bb5b..edd939f0368e 100644 --- a/tensorrt_llm/_torch/disaggregation/native/rank_info.py +++ b/tensorrt_llm/_torch/disaggregation/native/rank_info.py @@ -59,6 +59,10 @@ def from_kv_cache_manager( m = kv_cache_manager.mapping kvm = kv_cache_manager enable_attention_dp = m.enable_attention_dp + # Keep AttentionInfo on attention-free PP stages so it can still carry + # the attention-DP topology used by Mamba transfers. A zero head count + # means that this rank has no local attention cache; AttentionPolicy + # must not perform head-ratio arithmetic for such ranks. kv_heads_per_rank = next((h for h in kvm.num_kv_heads_per_layer if h > 0), 0) return cls( instance_name=instance_name, diff --git a/tests/unittest/disaggregated/test_peer.py b/tests/unittest/disaggregated/test_peer.py index 8ddfdcec3d42..00114949afb5 100644 --- a/tests/unittest/disaggregated/test_peer.py +++ b/tests/unittest/disaggregated/test_peer.py @@ -210,6 +210,37 @@ def test_no_overlap(): assert overlap.ranks == [] +@pytest.mark.parametrize( + ("self_heads", "peer_heads"), + [(0, 0), (0, 2), (2, 0)], +) +def test_attention_free_stage_has_no_head_duplication( + self_heads: int, + peer_heads: int, +) -> None: + self_ri = make_rankinfo( + "self", + tp_size=2, + tp_rank=0, + kv_heads_per_rank=self_heads, + layer_num_per_pp=[2], + ) + peer_ri = make_rankinfo( + "peer", + tp_size=2, + tp_rank=0, + kv_heads_per_rank=peer_heads, + layer_num_per_pp=[2], + ) + reg, peer_ri = _make_peer_registrar_and_peer_ri(self_ri, peer_ri) + + overlap = reg.get_peer_overlap(peer_ri, peer_dp_rank=0) + + assert overlap.duplicate_head_factor == 1 + assert overlap.peer_duplicate_head_factor == 1 + assert overlap.ranks == [0] + + def test_pp_ratio_peer_smaller(): self_ri = make_rankinfo( "self", diff --git a/tests/unittest/disaggregated/test_rank_info.py b/tests/unittest/disaggregated/test_rank_info.py index 1df741fb30c3..3284bbfc7776 100644 --- a/tests/unittest/disaggregated/test_rank_info.py +++ b/tests/unittest/disaggregated/test_rank_info.py @@ -16,10 +16,12 @@ from types import SimpleNamespace import numpy as np +import pytest from tensorrt_llm import bindings from tensorrt_llm._torch.disaggregation.native import rank_info as rank_info_module from tensorrt_llm._torch.disaggregation.native.auxiliary import AuxBufferMeta +from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import MambaPolicy from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo @@ -122,3 +124,36 @@ def test_from_kv_cache_manager_uses_first_nonzero_kv_head_count(monkeypatch) -> info = RankInfo.from_kv_cache_manager("ctx", manager, device_id=0) assert info.attention.kv_heads_per_rank == 8 + + +def test_from_kv_cache_manager_preserves_attention_dp_on_attention_free_stage( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(rank_info_module, "build_page_table_from_manager", lambda _: None) + mapping = SimpleNamespace( + rank=1, + tp_size=2, + tp_rank=1, + pp_size=1, + pp_rank=0, + dp_size=1, + cp_size=1, + cp_rank=0, + enable_attention_dp=True, + ) + manager = SimpleNamespace( + mapping=mapping, + num_kv_heads_per_layer=[0, 0], + pp_layers=[0, 1], + tokens_per_block=32, + head_dim=128, + dtype=bindings.DataType.HALF, + kv_factor=2, + ) + + info = RankInfo.from_kv_cache_manager("ctx", manager, device_id=0) + + assert info.attention is not None + assert info.attention.kv_heads_per_rank == 0 + assert info.attention.enable_attention_dp + assert MambaPolicy._mamba_tp(info) == (1, 0) From 9b5fca5d4ce137532c881169fa8126df00de770b Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Thu, 23 Jul 2026 05:10:29 +0800 Subject: [PATCH 15/22] [Agent fix] Correct V2 Mamba disagg docs (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- docs/source/features/kvcache.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index bebd98f33e84..9ca17132af8c 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -104,12 +104,12 @@ kv_cache_config: This retains snapshots after the first 128 tokens, at the end of the prompt, and before the final 32 prompt tokens. Positions outside a particular prompt -are ignored. Exact explicit boundaries currently require aggregated serving -with `MambaHybridCacheManagerV2`, `max_beam_width=1`, and no KV connector. -Hybrid Mamba models use the V1 C++ compatibility manager by default; select V2 -explicitly with `use_kv_cache_manager_v2: true`. V1 and the current -disaggregated-serving route support periodic snapshots only, while V2 does not -yet support disaggregated serving. +are ignored. Exact explicit boundaries currently require +`MambaHybridCacheManagerV2`, `max_beam_width=1`, and no KV connector. Hybrid +Mamba models use the V1 C++ compatibility manager by default; select V2 +explicitly with `use_kv_cache_manager_v2: true`. In disaggregated serving, V2 +Mamba requires the Python NIXL transceiver (`transceiver_runtime: PYTHON`); V1 +routes support periodic snapshots only. ### KV Cache Salting for Secure Reuse From 9e6fbd54a454785cb4dec9568ba6f6ad1d419125 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Thu, 23 Jul 2026 11:44:01 +0800 Subject: [PATCH 16/22] [Agent fix] Default hybrid Mamba models to KV cache V2 (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- docs/source/features/kvcache.md | 9 +++++---- tensorrt_llm/_torch/models/modeling_nemotron_h.py | 12 ++++++++---- tensorrt_llm/_torch/models/modeling_qwen3_5.py | 5 ++--- tensorrt_llm/_torch/models/modeling_qwen3_next.py | 14 +++++++++++--- tensorrt_llm/_torch/pyexecutor/_util.py | 10 +++++----- .../_torch/executor/test_mamba_cache_manager.py | 8 +++++++- 6 files changed, 38 insertions(+), 20 deletions(-) diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index 9ca17132af8c..f8b8afc14ddf 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -106,10 +106,11 @@ This retains snapshots after the first 128 tokens, at the end of the prompt, and before the final 32 prompt tokens. Positions outside a particular prompt are ignored. Exact explicit boundaries currently require `MambaHybridCacheManagerV2`, `max_beam_width=1`, and no KV connector. Hybrid -Mamba models use the V1 C++ compatibility manager by default; select V2 -explicitly with `use_kv_cache_manager_v2: true`. In disaggregated serving, V2 -Mamba requires the Python NIXL transceiver (`transceiver_runtime: PYTHON`); V1 -routes support periodic snapshots only. +Mamba models select V2 by default when +`use_kv_cache_manager_v2: auto`; set it to `false` to select the V1 C++ +compatibility manager. In disaggregated serving, V2 Mamba requires the Python +NIXL transceiver (`transceiver_runtime: PYTHON`); V1 routes support periodic +snapshots only. ### KV Cache Salting for Secure Reuse diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 8edb6f08e37c..8516a3dbdb1e 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -985,11 +985,15 @@ def load_weights(self, def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: """Model-specific defaults for NemotronH. - Disables block reuse due to SSM/hybrid architecture constraints. + Uses KV cache manager V2 for the hybrid state layout. Block reuse + remains opt-in because it also requires a Mamba snapshot policy. """ - # TODO: Remove enable_block_reuse=False once KV cache block reuse - # is supported for Mamba/SSM-based models - return {"kv_cache_config": {"enable_block_reuse": False}} + return { + "kv_cache_config": { + "enable_block_reuse": False, + "use_kv_cache_manager_v2": True, + } + } @classmethod def get_preferred_transceiver_runtime(cls, diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_5.py b/tensorrt_llm/_torch/models/modeling_qwen3_5.py index 4bd9de60c284..3a5c0612226c 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_5.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_5.py @@ -669,9 +669,8 @@ def get_model_defaults(cls, llm_args): # model class (this VLM wrapper), not on the inner decoder. Both # inner LMs (`Qwen3_5MoeForCausalLM` / `Qwen3_5ForCausalLM`) inherit # `Qwen3NextForCausalLM`'s defaults unchanged, so delegate to it to - # propagate `enable_block_reuse=False` — the hybrid Mamba/SSM path - # doesn't support KV-cache block reuse. Without this the VLM path - # would silently fall back to the global default (block reuse on). + # propagate the V2 manager selection and keep block reuse disabled + # until a recurrent-state snapshot policy is configured. return Qwen3NextForCausalLM.get_model_defaults(llm_args) @classmethod diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index c0041d9e54c7..da001279571c 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py @@ -991,9 +991,17 @@ def __init__( @classmethod def get_model_defaults(cls, llm_args: 'TorchLlmArgs') -> dict: - # TODO: Remove enable_block_reuse=False once KV cache block reuse - # is supported for Mamba/SSM-based models - return {"kv_cache_config": {"enable_block_reuse": False}} + """Use V2 for the hybrid state layout. + + Block reuse remains opt-in because it also requires a recurrent-state + snapshot policy. + """ + return { + "kv_cache_config": { + "enable_block_reuse": False, + "use_kv_cache_manager_v2": True, + } + } @classmethod def get_preferred_transceiver_runtime(cls, diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 9731e70a3e8c..4f00fde762c0 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -116,11 +116,11 @@ def get_kv_cache_manager_cls( Callers that don't care about disagg can omit ``is_disagg`` and get the unified-pool default. - V1 is the default for hybrid Mamba models. V2 is selected only by an - explicit ``kv_cache_config.use_kv_cache_manager_v2=True``. In - disaggregated serving, V2 additionally requires the Python transceiver - with the NIXL backend. Unsupported explicit V2 routes fail rather than - falling back to a different manager. + Model loading resolves ``use_kv_cache_manager_v2="auto"`` to V2 for + supported hybrid Mamba models. An explicit ``False`` selects a + compatibility manager. In disaggregated serving, V2 additionally requires + the Python transceiver with the NIXL backend. Unsupported V2 routes fail + rather than falling back to a different manager. Env-var overrides: * ``TRTLLM_USE_PY_MAMBA=1`` — Mixed manager in aggregated serving. diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index b6a2d36135b2..b7c0ce976012 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -62,6 +62,7 @@ from tensorrt_llm.llmapi.llm_utils import ( _resolve_kv_cache_manager_v2_auto, _resolve_transceiver_runtime_auto, + apply_model_defaults_to_llm_args, ) from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( @@ -570,7 +571,7 @@ def test_hybrid_cache_manager_factory_keeps_v1_disagg_route(monkeypatch, use_v2) ) -def test_hybrid_models_resolve_auto_to_python_transceiver(monkeypatch): +def test_hybrid_models_default_to_v2_and_python_transceiver(monkeypatch): from tensorrt_llm._torch.models.modeling_nemotron_h import NemotronHForCausalLM from tensorrt_llm._torch.models.modeling_qwen3_5 import Qwen3_5VLModel from tensorrt_llm._torch.models.modeling_qwen3_next import Qwen3NextForCausalLM @@ -588,7 +589,12 @@ def test_hybrid_models_resolve_auto_to_python_transceiver(monkeypatch): model="/tmp/dummy_model", cache_transceiver_config=CacheTransceiverConfig(backend="DEFAULT"), ) + model_defaults = model_cls.get_model_defaults(llm_args) + apply_model_defaults_to_llm_args(llm_args, model_defaults) + _resolve_kv_cache_manager_v2_auto(llm_args, model_defaults) _resolve_transceiver_runtime_auto(llm_args, model_cls) + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True + assert llm_args.kv_cache_config.enable_block_reuse is False assert llm_args.cache_transceiver_config.transceiver_runtime == "PYTHON" From b060d39d4ad57034b086aa7d5cbc89d8b515bc6f Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Thu, 23 Jul 2026 19:02:05 +0800 Subject: [PATCH 17/22] Report SSM snapshot reuse statistics (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 39 ++++++++- .../_torch/pyexecutor/kv_cache_stats.py | 65 +++++++++++++-- tensorrt_llm/metrics/collector.py | 10 ++- .../runtime/kv_cache_manager_v2/__init__.py | 2 + .../runtime/kv_cache_manager_v2/__init__.pyi | 15 ++++ .../kv_cache_manager_v2/_block_radix_tree.py | 7 +- .../kv_cache_manager_v2/_core/_kv_cache.py | 12 +++ .../_core/_kv_cache_manager.py | 30 ++++++- .../_core/_pending_stats.py | 40 ++++++++- .../runtime/kv_cache_manager_v2/_stats.py | 17 ++++ .../executor/test_kv_cache_manager_v2.py | 17 +++- .../executor/test_stats_serializer.py | 34 +++++++- .../test_kv_cache_manager_v2.py | 83 +++++++++++++++++++ .../test_kv_cache_stats_api.py | 16 ++++ tests/unittest/metrics/test_collector.py | 60 +++++++++++++- 15 files changed, 431 insertions(+), 16 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 728d1b17f937..8048ee689fa5 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -86,6 +86,8 @@ KVCacheV2IterationStatsReport, KVCacheV2LifeCycleIterationStats, KVCacheV2PoolGroupIterationStats, + KVCacheV2SsmLifeCycleIterationStats, + KVCacheV2SsmSnapshotIterationStats, ) from .llm_request import LlmRequest, LlmRequestState, SamplingConfig, get_draft_token_length from .resource_manager import ( @@ -2799,7 +2801,7 @@ def _build_pool_group_iteration_stats( ), ) - def _build_life_cycle_iteration_stats( + def _build_attention_life_cycle_iteration_stats( self, life_cycle_id: int, life_cycle_metadata, @@ -2810,6 +2812,7 @@ def _build_life_cycle_iteration_stats( reuse_delta, ) -> KVCacheV2LifeCycleIterationStats: pool_group_id, window_size, kind = life_cycle_metadata[life_cycle_id] + assert kind == "attention" return KVCacheV2LifeCycleIterationStats( life_cycle_id=life_cycle_id, pool_group_id=pool_group_id, @@ -2826,6 +2829,29 @@ def _build_life_cycle_iteration_stats( ), ) + @staticmethod + def _build_ssm_life_cycle_iteration_stats( + life_cycle_id: int, + life_cycle_metadata, + snapshot_delta, + ) -> KVCacheV2SsmLifeCycleIterationStats: + pool_group_id, window_size, kind = life_cycle_metadata[life_cycle_id] + assert kind == "ssm" + assert window_size is None + return KVCacheV2SsmLifeCycleIterationStats( + life_cycle_id=life_cycle_id, + pool_group_id=pool_group_id, + snapshot_stats=KVCacheV2SsmSnapshotIterationStats( + iter_snapshot_lookups=snapshot_delta.iter_snapshot_lookups, + iter_snapshot_hits=snapshot_delta.iter_snapshot_hits, + iter_snapshot_misses=snapshot_delta.iter_snapshot_misses, + iter_reused_tokens=snapshot_delta.iter_reused_tokens, + iter_unreused_tokens=snapshot_delta.iter_unreused_tokens, + iter_aligned_snapshot_hits=snapshot_delta.iter_aligned_snapshot_hits, + iter_unaligned_snapshot_hits=snapshot_delta.iter_unaligned_snapshot_hits, + ), + ) + def get_kv_cache_stats(self): kv_cache_stats = KvCacheStats() pool_group_stats = self._get_storage_statistics(GPU_LEVEL) @@ -2874,6 +2900,7 @@ def get_iteration_stats(self): pool_groups_by_window = self._storage_pool_groups_by_window() windows_by_pool_group = self._windows_by_pool_group(pool_groups_by_window) raw_iteration_stats = self.impl.get_and_reset_iteration_stats() + raw_ssm_snapshot_iteration_stats = self.impl.get_and_reset_ssm_snapshot_iteration_stats() primary_peak_stats = self._get_and_reset_iteration_peak_block_stats(GPU_LEVEL) num_cache_levels = len(self.impl.cache_tier_list) secondary_peak_stats_by_level = [ @@ -2928,7 +2955,7 @@ def get_iteration_stats(self): } stats_by_life_cycle = { - life_cycle_id: self._build_life_cycle_iteration_stats( + life_cycle_id: self._build_attention_life_cycle_iteration_stats( life_cycle_id, life_cycle_metadata, primary_stats, @@ -2939,6 +2966,14 @@ def get_iteration_stats(self): ) for life_cycle_id, reuse_delta in sorted(reuse_deltas_by_life_cycle.items()) } + for life_cycle_id, snapshot_delta in sorted(raw_ssm_snapshot_iteration_stats.items()): + assert int(life_cycle_id) not in stats_by_life_cycle + stats_by_life_cycle[int(life_cycle_id)] = self._build_ssm_life_cycle_iteration_stats( + int(life_cycle_id), + life_cycle_metadata, + snapshot_delta, + ) + stats_by_life_cycle = dict(sorted(stats_by_life_cycle.items())) return KVCacheV2IterationStatsReport( stats_by_window, stats_by_pool_group, stats_by_life_cycle diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_stats.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_stats.py index d9841ccbfe86..5ffa3742f3ef 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_stats.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_stats.py @@ -14,7 +14,7 @@ # limitations under the License. from dataclasses import dataclass, field -from typing import Any +from typing import Any, Literal KV_CACHE_ITERATION_STATS_REUSE_KEYS = ( "iterReusedBlocks", @@ -70,11 +70,39 @@ class KVCacheV2LifeCycleIterationStats: stats: Any +@dataclass(slots=True) +class KVCacheV2SsmSnapshotIterationStats: + iter_snapshot_lookups: int + iter_snapshot_hits: int + iter_snapshot_misses: int + iter_reused_tokens: int + iter_unreused_tokens: int + iter_aligned_snapshot_hits: int + iter_unaligned_snapshot_hits: int + + @property + def iter_snapshot_hit_rate(self) -> float: + if self.iter_snapshot_hits == 0 or self.iter_snapshot_lookups == 0: + return 0.0 + return self.iter_snapshot_hits / self.iter_snapshot_lookups + + +@dataclass(slots=True) +class KVCacheV2SsmLifeCycleIterationStats: + life_cycle_id: int + pool_group_id: int + snapshot_stats: KVCacheV2SsmSnapshotIterationStats + window_size: None = field(default=None, init=False) + kind: Literal["ssm"] = field(default="ssm", init=False) + + @dataclass(slots=True) class KVCacheV2IterationStatsReport: by_window_size: dict[int, Any] by_pool_group: dict[int, KVCacheV2PoolGroupIterationStats] - by_life_cycle: dict[int, KVCacheV2LifeCycleIterationStats] = field(default_factory=dict) + by_life_cycle: dict[ + int, KVCacheV2LifeCycleIterationStats | KVCacheV2SsmLifeCycleIterationStats + ] = field(default_factory=dict) def serialize_kv_cache_iteration_stats(stats, keys: tuple[str, ...] | None = None) -> dict: @@ -115,6 +143,21 @@ def serialize_kv_cache_iteration_stats(stats, keys: tuple[str, ...] | None = Non return {key: fields[key] for key in keys} +def serialize_ssm_snapshot_iteration_stats( + stats: KVCacheV2SsmSnapshotIterationStats, +) -> dict: + return { + "iterSnapshotLookups": stats.iter_snapshot_lookups, + "iterSnapshotHits": stats.iter_snapshot_hits, + "iterSnapshotMisses": stats.iter_snapshot_misses, + "iterSnapshotHitRate": stats.iter_snapshot_hit_rate, + "iterReusedTokens": stats.iter_reused_tokens, + "iterUnreusedTokens": stats.iter_unreused_tokens, + "iterAlignedSnapshotHits": stats.iter_aligned_snapshot_hits, + "iterUnalignedSnapshotHits": stats.iter_unaligned_snapshot_hits, + } + + def append_kv_cache_iteration_stats(stats_dict: dict, kv_iter_stats) -> None: if kv_iter_stats is None: return @@ -147,13 +190,21 @@ def append_kv_cache_iteration_stats(stats_dict: dict, kv_iter_stats) -> None: if not kv_iter_stats.by_life_cycle: return - stats_dict["kvCacheIterationStatsByLifecycle"] = { - str(life_cycle_id): { + stats_by_life_cycle = {} + for life_cycle_id, stats in kv_iter_stats.by_life_cycle.items(): + serialized = { "lifeCycleId": stats.life_cycle_id, "poolGroupId": stats.pool_group_id, "windowSize": stats.window_size, "kind": stats.kind, - **serialize_kv_cache_iteration_stats(stats.stats, KV_CACHE_ITERATION_STATS_REUSE_KEYS), } - for life_cycle_id, stats in kv_iter_stats.by_life_cycle.items() - } + if isinstance(stats, KVCacheV2SsmLifeCycleIterationStats): + serialized["snapshotStats"] = serialize_ssm_snapshot_iteration_stats( + stats.snapshot_stats + ) + else: + serialized.update( + serialize_kv_cache_iteration_stats(stats.stats, KV_CACHE_ITERATION_STATS_REUSE_KEYS) + ) + stats_by_life_cycle[str(life_cycle_id)] = serialized + stats_dict["kvCacheIterationStatsByLifecycle"] = stats_by_life_cycle diff --git a/tensorrt_llm/metrics/collector.py b/tensorrt_llm/metrics/collector.py index cd5a91123ddf..0586b46173ab 100644 --- a/tensorrt_llm/metrics/collector.py +++ b/tensorrt_llm/metrics/collector.py @@ -805,7 +805,15 @@ def log_iteration_stats(self, iteration_stats: dict) -> None: kv_iter_by_pool_group = iteration_stats.get( "kvCacheIterationStatsByPoolGroup") if kv_iter or kv_iter_by_lifecycle or kv_iter_by_pool_group: - reuse_stats = kv_iter_by_lifecycle or kv_iter or {} + # Prefer lifecycle-level attention stats when present. An SSM-only + # lifecycle report must not hide the legacy/window-level attention + # aggregate. Missing kind remains attention-compatible. + attention_lifecycle_stats = { + key: stats + for key, stats in (kv_iter_by_lifecycle or {}).items() + if stats.get("kind", "attention") == "attention" + } + reuse_stats = attention_lifecycle_stats or kv_iter or {} pool_group_stats = kv_iter_by_pool_group or kv_iter or {} total_secondary_max = 0 total_secondary_used = 0 diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 5557a1d2f32e..e96c62b5501c 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -78,6 +78,7 @@ _KV_CACHE_ITERATION_STATS_DELTA_FIELDS, KVCacheIterationStatsDelta, KVCacheStatsDelta, + SsmSnapshotIterationStatsDelta, ) from ._storage import BufferId # noqa: F401 from ._storage._config import CoalescedBuffer, SlotDesc, SlotDescVariant # noqa: F401 @@ -142,6 +143,7 @@ "SlotDesc", "SlotDescVariant", "SsmLayerConfig", + "SsmSnapshotIterationStatsDelta", "SwaScratchReuseConfig", "TokenId", "TokenIdExt", diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index acdcfeb1d608..dee3b8022d8a 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -93,6 +93,18 @@ class KVCacheIterationStatsDelta: iter_host_dropped_blocks: int = 0 iter_host_dropped_bytes: int = 0 +@dataclass(slots=True) +class SsmSnapshotIterationStatsDelta: + iter_snapshot_lookups: int = 0 + iter_snapshot_hits: int = 0 + iter_snapshot_misses: int = 0 + iter_reused_tokens: int = 0 + iter_unreused_tokens: int = 0 + iter_aligned_snapshot_hits: int = 0 + iter_unaligned_snapshot_hits: int = 0 + @property + def iter_snapshot_hit_rate(self) -> float: ... + @dataclass(slots=True, frozen=True) class PoolGroupPeakBlockStats: available: int @@ -486,6 +498,9 @@ class KVCacheManager: def get_quota(self, cache_level: CacheLevel) -> int: ... def get_committed_stats(self) -> KVCacheStatsDelta: ... def get_and_reset_iteration_stats(self) -> dict[LifeCycleId, KVCacheIterationStatsDelta]: ... + def get_and_reset_ssm_snapshot_iteration_stats( + self, + ) -> dict[LifeCycleId, SsmSnapshotIterationStatsDelta]: ... def get_and_reset_iteration_peak_block_stats( self, cache_level: CacheLevel ) -> Sequence[PoolGroupPeakBlockStats]: ... diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py index 4d55539b6b05..110912a07ddb 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py @@ -67,6 +67,7 @@ class ReuseMatch(NamedTuple): blocks: list["Block"] num_tokens: int + num_lookup_tokens: int # id_offset is usually vocab_size @@ -678,7 +679,11 @@ def match( matched = self._prune_match( list(self._match_token_path(reuse_scope, tokens, enable_partial_match)) ) - return ReuseMatch([block for block, _ in matched], self._num_matched_tokens(matched)) + return ReuseMatch( + [block for block, _ in matched], + self._num_matched_tokens(matched), + len(tokens), + ) def _check_sanity(self) -> bool: raise NotImplementedError( diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py index 9e20abbaeba2..d1aba2f540fd 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py @@ -412,6 +412,9 @@ def commit_pending_stats(self) -> KVCacheStatsDelta: self.manager.commit_stats( self._pending_stats.global_stats, self._pending_stats.iteration_stats_by_life_cycle ) + self.manager._commit_ssm_snapshot_iteration_stats( + self._pending_stats.ssm_snapshot_iteration_stats_by_life_cycle + ) request_stats = self._pending_stats.request_stats.copy() self._pending_stats.clear() self.manager.clear_stats_dirty(self.id) @@ -2081,6 +2084,15 @@ def _setup_for_reuse(self, match: ReuseMatch) -> None: ) snapshot_holder = unwrap_rawref(snapshot_ref).hold() self._ssm_blocks[DEFAULT_BEAM_INDEX][ssm_lc_id] = snapshot_holder + if should_record_stats and ssm_lc_id is not None: + changed = self._pending_stats.record_ssm_snapshot_lookup( + ssm_lc_id, + lookup_tokens=match.num_lookup_tokens, + reused_tokens=num_tokens, + tokens_per_block=tokens_per_block, + ) + if changed: + self.manager.mark_stats_dirty(self.id) self._num_committed_blocks = BlockOrdinal(len(self._committed_tokens) // tokens_per_block) for beam_indices in self._base_page_indices: for indices in beam_indices: diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py index 2d991b617db3..2492278bc9df 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py @@ -39,7 +39,7 @@ from .._config import DataRole, KVCacheManagerConfig from .._life_cycle_registry import LayerGroupId, LifeCycle, LifeCycleId, LifeCycleRegistry from .._page import Page, _PageHolder -from .._stats import KVCacheIterationStatsDelta, KVCacheStatsDelta +from .._stats import KVCacheIterationStatsDelta, KVCacheStatsDelta, SsmSnapshotIterationStatsDelta from .._storage._config import BufferId, SlotDesc, create_storage_config from .._storage._core import PoolGroupIndex, PoolIndex, SlotId from .._storage_manager import StorageManager @@ -212,6 +212,7 @@ class KVCacheManager: "_stats_enabled", "_committed_stats", "_iteration_stats_by_life_cycle", + "_ssm_snapshot_iteration_stats_by_life_cycle", "_iteration_peak_num_blocks_by_cache_level", "_dirty_stats_kv_cache_ids", "_stats_excluded_kv_cache_ids", @@ -241,6 +242,7 @@ class KVCacheManager: _stats_enabled: bool _committed_stats: KVCacheStatsDelta _iteration_stats_by_life_cycle: dict[LifeCycleId, KVCacheIterationStatsDelta] + _ssm_snapshot_iteration_stats_by_life_cycle: dict[LifeCycleId, SsmSnapshotIterationStatsDelta] _iteration_peak_num_blocks_by_cache_level: TypedIndexList[ CacheLevel, TypedIndexList[PoolGroupIndex, PoolGroupPeakBlockStats] ] @@ -284,6 +286,7 @@ def __init__( self._stats_enabled = config.enable_stats self._committed_stats = KVCacheStatsDelta() self._iteration_stats_by_life_cycle = {} + self._ssm_snapshot_iteration_stats_by_life_cycle = {} self._reset_iteration_peak_num_blocks() self._dirty_stats_kv_cache_ids = set() self._stats_excluded_kv_cache_ids = set() @@ -539,6 +542,31 @@ def get_and_reset_iteration_stats(self) -> dict[LifeCycleId, KVCacheIterationSta self._iteration_stats_by_life_cycle.clear() return stats + def _commit_ssm_snapshot_iteration_stats( + self, + iteration_stats_by_life_cycle: dict[LifeCycleId, SsmSnapshotIterationStatsDelta], + ) -> None: + if not self._stats_enabled: + return + for life_cycle, iteration_stats in iteration_stats_by_life_cycle.items(): + if iteration_stats.empty: + continue + destination = self._ssm_snapshot_iteration_stats_by_life_cycle.setdefault( + life_cycle, SsmSnapshotIterationStatsDelta() + ) + destination.add(iteration_stats) + + def get_and_reset_ssm_snapshot_iteration_stats( + self, + ) -> dict[LifeCycleId, SsmSnapshotIterationStatsDelta]: + stats = { + life_cycle: delta.copy() + for life_cycle, delta in self._ssm_snapshot_iteration_stats_by_life_cycle.items() + if not delta.empty + } + self._ssm_snapshot_iteration_stats_by_life_cycle.clear() + return stats + def get_and_reset_iteration_peak_block_stats( self, cache_level: CacheLevel ) -> TypedIndexList[PoolGroupIndex, PoolGroupPeakBlockStats]: diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_pending_stats.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_pending_stats.py index 94d0957537be..bf8318e1c726 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_pending_stats.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_pending_stats.py @@ -17,7 +17,7 @@ from .._common import BlockOrdinal from .._life_cycle_registry import LifeCycleId -from .._stats import KVCacheIterationStatsDelta, KVCacheStatsDelta +from .._stats import KVCacheIterationStatsDelta, KVCacheStatsDelta, SsmSnapshotIterationStatsDelta @dataclass(slots=True) @@ -49,6 +49,9 @@ class _PendingStats: iteration_stats_by_life_cycle: dict[LifeCycleId, KVCacheIterationStatsDelta] = field( default_factory=dict ) + ssm_snapshot_iteration_stats_by_life_cycle: dict[ + LifeCycleId, SsmSnapshotIterationStatsDelta + ] = field(default_factory=dict) allocation_segments: list[_PendingAllocationSegment] = field(default_factory=list) @property @@ -57,12 +60,14 @@ def empty(self) -> bool: self.request_stats.empty and self.global_stats.empty and not self.iteration_stats_by_life_cycle + and not self.ssm_snapshot_iteration_stats_by_life_cycle ) def clear(self) -> None: self.request_stats.clear() self.global_stats.clear() self.iteration_stats_by_life_cycle.clear() + self.ssm_snapshot_iteration_stats_by_life_cycle.clear() self.allocation_segments.clear() def add(self, delta: _PendingStatsDelta) -> bool: @@ -165,6 +170,39 @@ def record_reuse( ) ) + def record_ssm_snapshot_lookup( + self, + life_cycle: LifeCycleId, + *, + lookup_tokens: int, + reused_tokens: int, + tokens_per_block: int, + ) -> bool: + if lookup_tokens == 0: + return False + assert lookup_tokens > 0 + assert 0 <= reused_tokens <= lookup_tokens + assert tokens_per_block > 0 + + is_hit = reused_tokens > 0 + # Alignment describes the reusable snapshot boundary, not whether the + # state itself is complete. Every hit represents one complete SSM + # snapshot; token counters carry the benefit of that single lookup. + delta = SsmSnapshotIterationStatsDelta( + iter_snapshot_lookups=1, + iter_snapshot_hits=int(is_hit), + iter_snapshot_misses=int(not is_hit), + iter_reused_tokens=reused_tokens, + iter_unreused_tokens=lookup_tokens - reused_tokens, + iter_aligned_snapshot_hits=int(is_hit and reused_tokens % tokens_per_block == 0), + iter_unaligned_snapshot_hits=int(is_hit and reused_tokens % tokens_per_block != 0), + ) + pending = self.ssm_snapshot_iteration_stats_by_life_cycle.setdefault( + life_cycle, SsmSnapshotIterationStatsDelta() + ) + pending.add(delta) + return True + def subtract_allocation_range(self, block_begin: BlockOrdinal, block_end: BlockOrdinal) -> bool: if block_begin >= block_end or not self.allocation_segments: return False diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py index 292876ff4985..03d95d745df0 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py @@ -78,6 +78,23 @@ def iter_cache_hit_rate(self) -> float: return self.iter_reused_blocks / total +@dataclass(slots=True) +class SsmSnapshotIterationStatsDelta(_StatsDeltaMixin): + iter_snapshot_lookups: int = 0 + iter_snapshot_hits: int = 0 + iter_snapshot_misses: int = 0 + iter_reused_tokens: int = 0 + iter_unreused_tokens: int = 0 + iter_aligned_snapshot_hits: int = 0 + iter_unaligned_snapshot_hits: int = 0 + + @property + def iter_snapshot_hit_rate(self) -> float: + if self.iter_snapshot_hits == 0 or self.iter_snapshot_lookups == 0: + return 0.0 + return self.iter_snapshot_hits / self.iter_snapshot_lookups + + _KV_CACHE_ITERATION_STATS_DELTA_FIELDS = tuple( field.name for field in fields(KVCacheIterationStatsDelta) ) diff --git a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py index cf30c194f631..e45be30dcab0 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py @@ -513,11 +513,21 @@ def test_per_conversation_policy_ignores_overlapping_request( def test_iteration_stats_reports_physical_pool_groups_without_window_metadata() -> None: manager = object.__new__(KVCacheManagerV2) manager.enable_stats = True + snapshot_delta = SimpleNamespace( + iter_snapshot_lookups=2, + iter_snapshot_hits=1, + iter_snapshot_misses=1, + iter_reused_tokens=32, + iter_unreused_tokens=16, + iter_aligned_snapshot_hits=1, + iter_unaligned_snapshot_hits=0, + ) manager.impl = SimpleNamespace( cache_tier_list=[object()], get_and_reset_iteration_stats=lambda: {}, + get_and_reset_ssm_snapshot_iteration_stats=lambda: {3: snapshot_delta}, ) - manager._stats_life_cycle_metadata = lambda: {} + manager._stats_life_cycle_metadata = lambda: {3: (1, None, "ssm")} manager._storage_pool_groups_by_window = lambda: {} manager._get_and_reset_iteration_peak_block_stats = lambda _level: [None, None] manager._get_storage_statistics = lambda _level: [object(), object()] @@ -526,6 +536,11 @@ def test_iteration_stats_reports_physical_pool_groups_without_window_metadata() stats = manager.get_iteration_stats() assert stats.by_pool_group == {0: 0, 1: 1} + ssm_stats = stats.by_life_cycle[3] + assert ssm_stats.kind == "ssm" + assert ssm_stats.pool_group_id == 1 + assert ssm_stats.snapshot_stats.iter_snapshot_hit_rate == 0.5 + assert ssm_stats.snapshot_stats.iter_reused_tokens == 32 def test_disagg_role_mapper_kinds_default_to_indexed(): diff --git a/tests/unittest/executor/test_stats_serializer.py b/tests/unittest/executor/test_stats_serializer.py index bb061e4ec80d..b051947966ae 100644 --- a/tests/unittest/executor/test_stats_serializer.py +++ b/tests/unittest/executor/test_stats_serializer.py @@ -24,6 +24,8 @@ KVCacheV2IterationStatsReport, KVCacheV2LifeCycleIterationStats, KVCacheV2PoolGroupIterationStats, + KVCacheV2SsmLifeCycleIterationStats, + KVCacheV2SsmSnapshotIterationStats, ) from tensorrt_llm.executor.base_worker import BaseWorker @@ -387,7 +389,20 @@ def test_serializer_with_v2_pool_group_stats(self): window_size=16, kind="attention", stats=life_cycle_stats, - ) + ), + 4: KVCacheV2SsmLifeCycleIterationStats( + life_cycle_id=4, + pool_group_id=8, + snapshot_stats=KVCacheV2SsmSnapshotIterationStats( + iter_snapshot_lookups=4, + iter_snapshot_hits=3, + iter_snapshot_misses=1, + iter_reused_tokens=96, + iter_unreused_tokens=32, + iter_aligned_snapshot_hits=2, + iter_unaligned_snapshot_hits=1, + ), + ), }, ) @@ -425,6 +440,23 @@ def test_serializer_with_v2_pool_group_stats(self): assert life_cycle["iterReusedBlocks"] == 5 assert life_cycle["iterMissedBlocks"] == 3 assert "iterGenAllocBlocks" not in life_cycle + ssm_life_cycle = d["kvCacheIterationStatsByLifecycle"]["4"] + assert ssm_life_cycle == { + "lifeCycleId": 4, + "poolGroupId": 8, + "windowSize": None, + "kind": "ssm", + "snapshotStats": { + "iterSnapshotLookups": 4, + "iterSnapshotHits": 3, + "iterSnapshotMisses": 1, + "iterSnapshotHitRate": 0.75, + "iterReusedTokens": 96, + "iterUnreusedTokens": 32, + "iterAlignedSnapshotHits": 2, + "iterUnalignedSnapshotHits": 1, + }, + } def test_v2_peak_block_stats_reset_tracks_interval_peak(self): """Peak block stats should cover the interval since the previous reset.""" diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index 19e23ffb36a5..7dba91291797 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -1989,6 +1989,89 @@ def test_no_reuse_with_ssm(self) -> None: kv_cache.resume(stream) kv_cache.close() + @parameterized.expand( + [ + ("miss", None, 48, False, (1, 0, 1, 0, 48, 0, 0)), + ("aligned_hit", 32, 48, False, (1, 1, 0, 32, 16, 1, 0)), + ("unaligned_hit", 48, 64, True, (1, 1, 0, 48, 16, 0, 1)), + ] + ) + def test_ssm_snapshot_iteration_stats( + self, + _name: str, + snapshot_length: int | None, + lookup_length: int, + enable_partial_reuse: bool, + expected: tuple[int, int, int, int, int, int, int], + ) -> None: + tokens_per_block = 32 + cfg = self._make_ssm_config( + tokens_per_block=tokens_per_block, + enable_partial_reuse=enable_partial_reuse, + ) + self.manager = KVCacheManager(cfg) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + prompt = [self.next_token() for _ in range(lookup_length)] + + if snapshot_length is not None: + seed = self.manager.create_kv_cache() + seed.resume(stream) + seed.capacity = snapshot_length + seed.history_length = snapshot_length + seed.commit(prompt[:snapshot_length], is_end=True) + seed.close() + + reused = self.manager.create_kv_cache( + input_tokens=prompt, + id=101, + # This is only a sizing hint; lookup telemetry must use the + # actual input_tokens length. + expected_prompt_length=lookup_length + 17, + ) + self.assertEqual(reused.num_committed_tokens, expected[3]) + self.assertEqual(self.manager.get_dirty_stats_kv_cache_ids(), {101}) + reused.commit_pending_stats() + self.assertEqual(self.manager.get_dirty_stats_kv_cache_ids(), set()) + + assert self.manager._life_cycles.ssm_life_cycle_id is not None + ssm_life_cycle_id = self.manager._life_cycles.ssm_life_cycle_id + snapshot_stats = self.manager.get_and_reset_ssm_snapshot_iteration_stats() + self.assertEqual(set(snapshot_stats), {ssm_life_cycle_id}) + stats = snapshot_stats[ssm_life_cycle_id] + self.assertEqual( + ( + stats.iter_snapshot_lookups, + stats.iter_snapshot_hits, + stats.iter_snapshot_misses, + stats.iter_reused_tokens, + stats.iter_unreused_tokens, + stats.iter_aligned_snapshot_hits, + stats.iter_unaligned_snapshot_hits, + ), + expected, + ) + self.assertEqual(stats.iter_snapshot_hit_rate, expected[1] / expected[0]) + self.assertEqual(self.manager.get_and_reset_ssm_snapshot_iteration_stats(), {}) + + reused.resume(stream) + reused.close() + + def test_discard_ssm_snapshot_stats_clears_dirty_state(self) -> None: + cfg = self._make_ssm_config() + self.manager = KVCacheManager(cfg) + tokens = [self.next_token() for _ in range(16)] + + kv_cache = self.manager.create_kv_cache(input_tokens=tokens, id=101) + self.assertEqual(self.manager.get_dirty_stats_kv_cache_ids(), {101}) + kv_cache.discard_pending_stats() + + self.assertEqual(self.manager.get_dirty_stats_kv_cache_ids(), set()) + self.assertEqual(self.manager.get_and_reset_ssm_snapshot_iteration_stats(), {}) + stream_holder = CachedCudaStream() + kv_cache.resume(cast(CudaStream, stream_holder.handle)) + kv_cache.close() + def test_ssm(self) -> None: """Inference with SSM layer: prefill 63 tokens, decode 52 tokens.""" cfg = self._make_ssm_config() diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_api.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_api.py index fa43bf2d0992..ea25d28e541c 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_api.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_api.py @@ -26,6 +26,7 @@ KVCacheManager, KVCacheManagerConfig, KVCacheStatsDelta, + SsmSnapshotIterationStatsDelta, ) pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -63,6 +64,20 @@ def test_stats_delta_arithmetic() -> None: assert iteration.empty assert iteration.iter_cache_hit_rate == 0.0 + snapshot = SsmSnapshotIterationStatsDelta( + iter_snapshot_lookups=4, + iter_snapshot_hits=3, + iter_snapshot_misses=1, + iter_reused_tokens=96, + iter_unreused_tokens=32, + iter_aligned_snapshot_hits=2, + iter_unaligned_snapshot_hits=1, + ) + assert snapshot.iter_snapshot_hit_rate == 0.75 + snapshot.clear() + assert snapshot.empty + assert snapshot.iter_snapshot_hit_rate == 0.0 + @pytest.mark.parametrize("enable_stats", [False, True]) def test_manager_stats_config_and_api(enable_stats: bool) -> None: @@ -72,6 +87,7 @@ def test_manager_stats_config_and_api(enable_stats: bool) -> None: assert manager.init_config.enable_stats is enable_stats assert manager.get_committed_stats() == KVCacheStatsDelta() assert manager.get_and_reset_iteration_stats() == {} + assert manager.get_and_reset_ssm_snapshot_iteration_stats() == {} peak_stats = manager.get_and_reset_iteration_peak_block_stats(GPU_LEVEL) assert len(peak_stats) == 1 assert peak_stats[0].available >= 0 diff --git a/tests/unittest/metrics/test_collector.py b/tests/unittest/metrics/test_collector.py index e66975ca6f8e..2727a36fb73e 100644 --- a/tests/unittest/metrics/test_collector.py +++ b/tests/unittest/metrics/test_collector.py @@ -814,7 +814,20 @@ def test_v2_lifecycle_and_pool_group_stats_are_aggregated(self): "iterFullReusedBlocks": 4, "iterPartialReusedBlocks": 1, "iterMissedBlocks": 3, - } + }, + "1": { + "kind": "ssm", + "snapshotStats": { + "iterSnapshotLookups": 2, + "iterSnapshotHits": 1, + "iterSnapshotMisses": 1, + "iterSnapshotHitRate": 0.5, + "iterReusedTokens": 128, + "iterUnreusedTokens": 64, + "iterAlignedSnapshotHits": 1, + "iterUnalignedSnapshotHits": 0, + }, + }, }, "kvCacheIterationStatsByPoolGroup": { "0": { @@ -829,6 +842,7 @@ def test_v2_lifecycle_and_pool_group_stats_are_aggregated(self): } before_reused = _get_counter_value(collector, "kv_cache_iter_reused_blocks") + before_missed = _get_counter_value(collector, "kv_cache_iter_missed_blocks") before_gen_alloc = _get_counter_value(collector, "kv_cache_gen_alloc_blocks_total") before_onboard = _get_counter_value(collector, "kv_cache_onboard_bytes_total") @@ -839,6 +853,9 @@ def test_v2_lifecycle_and_pool_group_stats_are_aggregated(self): assert _get_counter_value( collector, "kv_cache_iter_reused_blocks" ) - before_reused == pytest.approx(5) + assert _get_counter_value( + collector, "kv_cache_iter_missed_blocks" + ) - before_missed == pytest.approx(3) assert _get_counter_value( collector, "kv_cache_gen_alloc_blocks_total" ) - before_gen_alloc == pytest.approx(2) @@ -846,6 +863,47 @@ def test_v2_lifecycle_and_pool_group_stats_are_aggregated(self): collector, "kv_cache_onboard_bytes_total" ) - before_onboard == pytest.approx(4096) + def test_v2_ssm_only_lifecycle_falls_back_to_attention_window_stats(self): + """An SSM lifecycle entry must not hide attention's window aggregate.""" + collector = _make_kv_iter_collector() + stats = { + "kvCacheIterationStats": { + "16": { + "iterReusedBlocks": 2, + "iterFullReusedBlocks": 2, + "iterPartialReusedBlocks": 0, + "iterMissedBlocks": 2, + } + }, + "kvCacheIterationStatsByLifecycle": { + "1": { + "kind": "ssm", + "snapshotStats": { + "iterSnapshotLookups": 1, + "iterSnapshotHits": 1, + "iterSnapshotMisses": 0, + "iterSnapshotHitRate": 1.0, + "iterReusedTokens": 32, + "iterUnreusedTokens": 0, + "iterAlignedSnapshotHits": 1, + "iterUnalignedSnapshotHits": 0, + }, + } + }, + } + before_reused = _get_counter_value(collector, "kv_cache_iter_reused_blocks") + before_missed = _get_counter_value(collector, "kv_cache_iter_missed_blocks") + + collector.log_iteration_stats(stats) + + assert _get_gauge_value(collector, "kv_cache_iter_reuse_rate") == pytest.approx(0.5) + assert _get_counter_value( + collector, "kv_cache_iter_reused_blocks" + ) - before_reused == pytest.approx(2) + assert _get_counter_value( + collector, "kv_cache_iter_missed_blocks" + ) - before_missed == pytest.approx(2) + def test_multiple_windows_aggregated(self): """Stats from multiple window sizes should be summed.""" collector = _make_kv_iter_collector() From 6b529d4d82f905ed87c5d425cc432778da86d1ed Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Thu, 23 Jul 2026 20:42:28 +0800 Subject: [PATCH 18/22] Use default V2 for MTP accuracy tests (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../defs/accuracy/test_llm_api_pytorch.py | 20 +++---------------- .../test_lists/qa/llm_function_core.txt | 6 ++---- .../test_lists/qa/llm_function_rtx6k.txt | 3 +-- .../test_lists/test-db/l0_b200.yml | 3 +-- .../test_lists/test-db/l0_dgx_b200.yml | 3 +-- tests/integration/test_lists/waives.txt | 6 ++---- 6 files changed, 10 insertions(+), 31 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 1689bcf3adde..a412513137bf 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -6241,15 +6241,9 @@ def test_bf16(self, moe_backend, tp_size, mocker): extra_acc_spec=extra_acc_spec, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) - @pytest.mark.parametrize( - "v2_kv_cache", - [False, True], - ids=["v1_kv_cache", "v2_kv_cache"], - ) - def test_bf16_mtp(self, v2_kv_cache, mocker): + def test_bf16_mtp(self, mocker): kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.8, - enable_block_reuse=False, - use_kv_cache_manager_v2=v2_kv_cache) + enable_block_reuse=False) cuda_graph_config = CudaGraphConfig(enable_padding=True, max_batch_size=32) @@ -6268,7 +6262,6 @@ def test_bf16_mtp(self, v2_kv_cache, mocker): kv_cache_config=kv_cache_config, cuda_graph_config=cuda_graph_config, speculative_config=mtp_config) as llm: - assert llm.args.kv_cache_config.use_kv_cache_manager_v2 is v2_kv_cache mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN) task = GSM8K(self.MODEL_NAME) @@ -7171,12 +7164,7 @@ def test_nvfp4_parallelism(self, tp_size, ep_size, pp_size, attention_dp): @skip_pre_blackwell @pytest.mark.skip_less_mpi_world_size(8) - @pytest.mark.parametrize( - "v2_kv_cache", - [False, True], - ids=["v1_kv_cache", "v2_kv_cache"], - ) - def test_nvfp4_8gpus_mtp(self, v2_kv_cache): + def test_nvfp4_8gpus_mtp(self): # Test MTP (Multi-Token Prediction) accuracy with nvfp4-fp8kv model. # This test uses MTP with max_draft_len=3 and one_model mode. mtp_config = MTPDecodingConfig( @@ -7190,7 +7178,6 @@ def test_nvfp4_8gpus_mtp(self, v2_kv_cache): enable_block_reuse=True, mamba_ssm_cache_dtype="float16", free_gpu_memory_fraction=0.5, - use_kv_cache_manager_v2=v2_kv_cache, ), max_batch_size=32, tensor_parallel_size=8, @@ -7203,7 +7190,6 @@ def test_nvfp4_8gpus_mtp(self, v2_kv_cache): moe_config=MoeConfig(backend="CUTLASS"), decoding_config=mtp_config, ) as llm: - assert llm.args.kv_cache_config.use_kv_cache_manager_v2 is v2_kv_cache task = MMLU(self.MODEL_NAME) task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index a1e7cac4b001..d675ad632e6c 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -681,8 +681,7 @@ accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_online_e accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_static_eplb[moe_backend=CUTLASS] accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_static_eplb[moe_backend=TRTLLM] accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_static_eplb[moe_backend=CUTEDSL] -accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp[v1_kv_cache] -accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp[v2_kv_cache] +accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp_custom_op accuracy/test_llm_api_pytorch.py::TestNemotronV3Nano::test_nvfp4_marlin_multi_gpus[tp_size=8] accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_marlin_multi_gpus[tp_size=8] @@ -703,8 +702,7 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4[latency_moe_tr accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4[latency_moe_trtllm_attention_dp] accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4_4gpus[latency_moe_cutlass] accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4_4gpus[latency_moe_trtllm_eagle3] -accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v1_kv_cache] -accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v2_kv_cache] +accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_dummy_load_format accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_fp8[latency-torch_compile=False] accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_fp8[latency-torch_compile=True] diff --git a/tests/integration/test_lists/qa/llm_function_rtx6k.txt b/tests/integration/test_lists/qa/llm_function_rtx6k.txt index d16dc73c5745..4260a5328d77 100644 --- a/tests/integration/test_lists/qa/llm_function_rtx6k.txt +++ b/tests/integration/test_lists/qa/llm_function_rtx6k.txt @@ -173,8 +173,7 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[latency_moe_cutl accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[latency_moe_cutlass-torch_compile=True] accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp1-CUTLASS] accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp1-TRTLLM] -accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v1_kv_cache] -accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v2_kv_cache] +accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8[enable_block_reuse=False] accuracy/test_llm_api_pytorch.py::TestQwen3_5_397B_A17B::test_nvfp4[adp4_cutedsl] accuracy/test_llm_api_pytorch.py::TestQwen3_5_397B_A17B::test_nvfp4[adp4_trtllm] diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 8a6892bb4491..d812dcf5d84e 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -61,8 +61,7 @@ l0_b200: - accuracy/test_llm_api_pytorch_multimodal.py::TestNanoV3Omni::test_auto_dtype[fp8_mmmu_encoder_cuda_graph] - accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[qwen3vl_2b_instruct] - accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[nemotron_nano_v3_omni_nvfp4] - - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v1_kv_cache] - - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v2_kv_cache] + - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp - disaggregated/test_workers.py::test_workers_kv_cache_aware_router_eviction[TinyLlama-1.1B-Chat-v1.0] # nvbugs 5300551 - test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-NVFP4-nvfp4-quantized/Meta-Llama-3.1-8B] - test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-FP8-llama-3.1-model/Llama-3.1-8B-Instruct-FP8] diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index f006349d75d2..0a8988c17edc 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -169,8 +169,7 @@ l0_dgx_b200: - accuracy/test_disaggregated_serving.py::TestKimiK25::test_nvfp4 TIMEOUT (180) - accuracy/test_llm_api_pytorch.py::TestKimiK25::test_nvfp4[tp8] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4ProDSpark::test_gsm8k_dep8_megamoe_deepgemm TIMEOUT (240) - - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp[v1_kv_cache] TIMEOUT (60) - - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp[v2_kv_cache] TIMEOUT (60) + - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_8gpus[attention_dp_on-cutedsl] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_8gpus[attention_dp_off-trtllm] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=0-block_reuse=False-use_py_transceiver=True] TIMEOUT (60) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 7274098b7aa5..5f4e9ace70f0 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -277,8 +277,7 @@ full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_2gpus[ep2-t full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_2gpus[tp2-trtllm-auto] SKIP (https://nvbugs/6273846) full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_2gpus[tp2-trtllm-fp8] SKIP (https://nvbugs/6273846) full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp1-CUTLASS] SKIP (https://nvbugs/6273850) -full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v1_kv_cache] SKIP (https://nvbugs/6275856) -full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v2_kv_cache] SKIP (https://nvbugs/6275856) +full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp SKIP (https://nvbugs/6275856) full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8[enable_block_reuse=False] SKIP (https://nvbugs/6313076) full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_dflash SKIP (https://nvbugs/6273850) full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_fp8 SKIP (https://nvbugs/6273850) @@ -310,8 +309,7 @@ full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::Tes full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_2gpus[tp2-trtllm-auto] SKIP (https://nvbugs/6273846) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_2gpus[tp2-trtllm-fp8] SKIP (https://nvbugs/6273846) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp1-CUTLASS] SKIP (https://nvbugs/6273850) -full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v1_kv_cache] SKIP (https://nvbugs/6275856) -full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[v2_kv_cache] SKIP (https://nvbugs/6275856) +full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp SKIP (https://nvbugs/6275856) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8[enable_block_reuse=False] SKIP (https://nvbugs/6313076) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16 SKIP (https://nvbugs/6273850) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_dflash SKIP (https://nvbugs/6273850) From 615d124c3d3d36ec01f90220266feab2c3b14f3e Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Thu, 23 Jul 2026 21:17:38 +0800 Subject: [PATCH 19/22] [None][fix] Address pre-merge review findings (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_torch/disaggregation/transceiver.py | 4 +- tensorrt_llm/commands/serve.py | 12 ++++-- tensorrt_llm/llmapi/llm_args.py | 14 ++++++- .../executor/test_mamba_cache_manager.py | 23 +++++++++++ tests/unittest/llmapi/test_llm_args.py | 39 +++++++++++++++++-- 5 files changed, 82 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index bd2dbdbd44a3..d0bfcebe1095 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -237,7 +237,9 @@ def _create_kv_slice(self, req: LlmRequest) -> KVSlice: mamba_state_index = None if isinstance(self._kv_cache_manager, MambaHybridCacheManagerV2): if self._kv_cache_manager.local_num_mamba_layers > 0: - mamba_state_index = self._kv_cache_manager.get_state_indices([req.py_request_id])[0] + mamba_state_index = self._kv_cache_manager._request_id_to_state_index[ + req.py_request_id + ] elif isinstance(self._kv_cache_manager, MambaHybridCacheManager): mamba_state_index = self._kv_cache_manager.mamba_cache_index[req.py_request_id] diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 8766a2f260ee..c419ffdf263c 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -1387,7 +1387,9 @@ def _serve_llm(): if extra_llm_api_options is not None: with open(extra_llm_api_options, 'r') as f: llm_args_extra_dict = yaml.safe_load(f) - if not isinstance(llm_args_extra_dict, dict): + if llm_args_extra_dict is None: + llm_args_extra_dict = {} + elif not isinstance(llm_args_extra_dict, dict): raise ValueError("Configuration file root must be a mapping.") extra_allow_request_chat_template = _pop_bool_config_option( llm_args_extra_dict, "allow_request_chat_template") @@ -1624,7 +1626,9 @@ def serve_encoder(model: str, host: str, port: int, log_level: str, if extra_encoder_options is not None: with open(extra_encoder_options, 'r') as f: encoder_args_extra_dict = yaml.safe_load(f) - if not isinstance(encoder_args_extra_dict, dict): + if encoder_args_extra_dict is None: + encoder_args_extra_dict = {} + elif not isinstance(encoder_args_extra_dict, dict): raise ValueError("Configuration file root must be a mapping.") extra_allow_request_chat_template = _pop_bool_config_option( encoder_args_extra_dict, "allow_request_chat_template") @@ -1743,7 +1747,9 @@ def serve_embedding( if extra_llm_api_options is not None: with open(extra_llm_api_options, 'r') as f: extra_dict = yaml.safe_load(f) - if not isinstance(extra_dict, dict): + if extra_dict is None: + extra_dict = {} + elif not isinstance(extra_dict, dict): raise ValueError("Configuration file root must be a mapping.") llm_args = update_llm_args_with_extra_dict( llm_args, extra_dict, explicit_cli_keys=explicit_cli_keys) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 3c59cc15368b..bed46d766606 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3939,6 +3939,12 @@ def disable_periodic_mamba_snapshots_for_conversations( """Use only explicit stable boundaries for conversation reuse.""" if (self.block_reuse_policy == "per_conversation" and self.mamba_state_config.periodic_snapshot_interval != 0): + interval = self.mamba_state_config.periodic_snapshot_interval + logger.warning( + f"'kv_cache_config.mamba_state_config.periodic_snapshot_interval={interval}' " + "is ignored because " + "'kv_cache_config.block_reuse_policy=per_conversation' disables " + "periodic Mamba snapshots; setting it to 0.") self.mamba_state_config = self.mamba_state_config.model_copy( update={"periodic_snapshot_interval": 0}) return self @@ -4539,7 +4545,9 @@ def speculative_model(self) -> Optional[Union[str, Path]]: def from_yaml(cls, yaml_path: Union[str, Path]): with open(yaml_path, "r") as f: config_dict = yaml.safe_load(f) - if not isinstance(config_dict, dict): + if config_dict is None: + config_dict = {} + elif not isinstance(config_dict, dict): raise ValueError("Configuration file root must be a mapping.") return cls(**config_dict) @@ -6084,7 +6092,9 @@ def update_llm_args_with_extra_options( if extra_llm_api_options is not None: with open(extra_llm_api_options, 'r') as f: llm_args_dict = yaml.safe_load(f) - if not isinstance(llm_args_dict, dict): + if llm_args_dict is None: + llm_args_dict = {} + elif not isinstance(llm_args_dict, dict): raise ValueError("Configuration file root must be a mapping.") llm_args = update_llm_args_with_extra_dict( llm_args, diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index b7c0ce976012..1d0dbe3ad7d7 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -616,6 +616,29 @@ def test_v2_disagg_slice_skips_state_index_on_mamba_free_pp_rank(): assert kv_slice.mamba_state_index is None +def test_v2_disagg_slice_reads_state_index_without_refreshing_batch_mask(): + manager = object.__new__(MambaHybridCacheManagerV2) + manager.local_num_mamba_layers = 1 + manager._request_id_to_state_index = {123: 7} + manager.get_state_indices = MagicMock( + side_effect=AssertionError("state-index lookup must not refresh the dummy mask") + ) + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._kv_cache_manager = manager + transceiver._reuse_adapter = SimpleNamespace(tokens_per_block=32) + transceiver._page_table = SimpleNamespace(layer_groups=[]) + request = SimpleNamespace( + is_generation_only_request=lambda: False, + prompt_len=0, + py_request_id=123, + ) + + kv_slice = transceiver._create_kv_slice(request) + + assert kv_slice.mamba_state_index == 7 + manager.get_state_indices.assert_not_called() + + @pytest.mark.parametrize( "max_beam_width, has_connector, expected", [ diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index c06e50dd62c7..985944a96ed3 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -54,7 +54,8 @@ StrictBaseModel, TorchCompileConfig, TorchLlmArgs, UserProvidedDecodingConfig, - update_llm_args_with_extra_dict) + update_llm_args_with_extra_dict, + update_llm_args_with_extra_options) # fmt: on from tensorrt_llm.llmapi.llm_utils import (_resolve_kv_cache_manager_v2_auto, _resolve_transceiver_runtime_auto, @@ -198,12 +199,11 @@ def test_from_yaml_migrates_legacy_mamba_interval(self, tmp_path): assert llm_args.kv_cache_config.mamba_state_config.periodic_snapshot_interval == 64 - def test_from_yaml_rejects_empty_file(self, tmp_path): + def test_from_yaml_empty_file_reports_missing_model(self, tmp_path): yaml_path = tmp_path / "empty.yaml" yaml_path.write_text("", encoding="utf-8") - with pytest.raises(ValueError, - match="Configuration file root must be a mapping"): + with pytest.raises(ValidationError, match="model"): TorchLlmArgs.from_yaml(yaml_path) def test_from_yaml_rejects_non_mapping_root(self, tmp_path): @@ -807,6 +807,37 @@ def test_KvCacheConfig_migrates_deprecated_mamba_interval(monkeypatch): assert llm_args.kv_cache_config.mamba_state_config.periodic_snapshot_interval == 32 +def test_KvCacheConfig_warns_when_disabling_periodic_conversation_snapshots( + monkeypatch): + warnings_seen = [] + monkeypatch.setattr(llm_args_mod.logger, "warning", + lambda message: warnings_seen.append(message)) + + config = KvCacheConfig( + block_reuse_policy="per_conversation", + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=64), + ) + + assert config.mamba_state_config.periodic_snapshot_interval == 0 + assert len(warnings_seen) == 1 + assert "periodic_snapshot_interval=64" in warnings_seen[0] + assert "block_reuse_policy=per_conversation" in warnings_seen[0] + assert "setting it to 0" in warnings_seen[0] + + warnings_seen.clear() + KvCacheConfig(block_reuse_policy="per_conversation") + assert warnings_seen == [] + + +def test_update_llm_args_with_empty_options_file(tmp_path): + yaml_path = tmp_path / "empty.yaml" + yaml_path.write_text("", encoding="utf-8") + llm_args = {"model": "dummy"} + + assert update_llm_args_with_extra_options(llm_args, + str(yaml_path)) == llm_args + + def test_config_file_merge_migrates_legacy_mamba_interval_without_mutating_input( ): yaml_dict = { From d1ed15f9eca75e68717030d3e611bffecae6c525 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Thu, 23 Jul 2026 22:05:58 +0800 Subject: [PATCH 20/22] [Agent fix] Simplify V2 Mamba cache changes (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_torch/pyexecutor/mamba_cache_manager.py | 136 ++++-------- .../test_kv_cache_v2_capacity_only.py | 23 +- .../executor/test_mamba_cache_manager.py | 196 +++++------------- .../unittest/disaggregated/test_extractor.py | 10 - .../disaggregated/test_mamba_transfer.py | 86 +------- .../test_kv_cache_manager_v2.py | 5 - tests/unittest/llmapi/test_llm_args.py | 39 +--- 7 files changed, 106 insertions(+), 389 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 46fe3feea55c..425d26fc969d 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -45,7 +45,6 @@ from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import (DEFAULT_BEAM_INDEX, - AttentionLayerConfig, BatchDesc, BufferConfig, DataRole, KVCacheDesc) from tensorrt_llm.runtime.kv_cache_manager_v2 import \ @@ -1658,7 +1657,6 @@ def _estimate_mamba_hybrid_cache_cost( *, max_batch_size: int, kv_cache_config: KvCacheConfig, - num_layers: Optional[int], tokens_per_block: int, max_seq_len: Optional[int], num_reserved_dummy_slots: int, @@ -1668,7 +1666,6 @@ def _estimate_mamba_hybrid_cache_cost( use_separate_draft_kv_cache: bool = False, **kwargs, ) -> Tuple[int, int]: - del num_layers spec_config = kwargs.get("spec_config") params, local_mamba_layers, local_attention_layers = ( _get_local_mamba_cache_layout( @@ -1896,10 +1893,10 @@ def __init__( # Section dims are derived from mHiddenSize and mNGroups*mDState in C++; # we just need to tell C++ which ordering to use. self._rnn_conv_section_layout = model_type # "nemotron_hybrid" or "qwen3_next" - self.ssm_count = math.prod(self.ssm_state_shape) - self.conv_count = math.prod(self.conv_state_shape) - self.ssm_bytes = self.ssm_count * self.ssm_state_dtype.itemsize - self.conv_bytes = self.conv_count * self.conv_state_dtype.itemsize + self.ssm_bytes = (math.prod(self.ssm_state_shape) * + self.ssm_state_dtype.itemsize) + self.conv_bytes = (math.prod(self.conv_state_shape) * + self.conv_state_dtype.itemsize) total_bytes = self.ssm_bytes + self.conv_bytes if total_bytes % self.ssm_state_dtype.itemsize != 0: @@ -2023,10 +2020,6 @@ def __init__( device="cpu") self._request_id_to_state_index = {} self._request_id_to_is_dummy = {} - # Batch-order mask aligned with state_indices; duplicate dummy request - # IDs mark every batch row even when they share one cache slot. - self._dummy_request_mask = None - self._dummy_request_mask_host = None self.kv_cache_config = kv_cache_config self.is_estimating_kv_cache = is_estimating_kv_cache @@ -2068,7 +2061,6 @@ def get_cache_size_per_token( mapping, max_batch_size=max_batch_size, kv_cache_config=kv_cache_config, - num_layers=num_layers, tokens_per_block=tokens_per_block, max_seq_len=max_seq_len, num_reserved_dummy_slots=1, @@ -2549,13 +2541,11 @@ def __init__( self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding self._seed_rank_offset = _mamba_rank_offset(mapping) self._seed_request_counter = 0 - self._num_cuda_graph_padding_dummy_slots = ( + num_cuda_graph_padding_dummy_slots = ( _get_num_cuda_graph_padding_dummy_slots(spec_config, max_batch_size)) - self._num_attention_dp_dummy_slots = int(mapping.enable_attention_dp) - self._num_reserved_dummy_slots = ( - self._num_cuda_graph_padding_dummy_slots + - self._num_attention_dp_dummy_slots) + self._num_reserved_dummy_slots = (num_cuda_graph_padding_dummy_slots + + int(mapping.enable_attention_dp)) self.ssm_state_dtype = (mamba_ssm_cache_dtype if mamba_ssm_cache_dtype is not None else mamba_cache_dtype) self.conv_state_dtype = mamba_cache_dtype @@ -2611,10 +2601,10 @@ def __init__( grouped_state_dim_local, d_inner_local, ] - self.ssm_count = math.prod(self.ssm_state_shape) - self.conv_count = math.prod(self.conv_state_shape) - self.ssm_bytes = self.ssm_count * self.ssm_state_dtype.itemsize - self.conv_bytes = self.conv_count * self.conv_state_dtype.itemsize + self.ssm_bytes = (math.prod(self.ssm_state_shape) * + self.ssm_state_dtype.itemsize) + self.conv_bytes = (math.prod(self.conv_state_shape) * + self.conv_state_dtype.itemsize) else: logger.info( "No local mamba layers for this rank, skipping mamba state views" @@ -2623,8 +2613,6 @@ def __init__( self.conv_state_shape = [] self.ssm_state_shape = [] self.conv_section_dims = [] - self.ssm_count = 0 - self.conv_count = 0 self.ssm_bytes = 0 self.conv_bytes = 0 @@ -2675,9 +2663,6 @@ def __init__( layer_id: offset for offset, layer_id in enumerate(self.mamba_pp_layers) } - self.mamba_local_layer_ids = [ - self.layer_offsets[layer_id] for layer_id in self.mamba_pp_layers - ] self._request_id_to_state_index = {} self._request_id_to_is_dummy = {} @@ -2691,7 +2676,8 @@ def __init__( pin_memory=prefer_pinned()) if self.local_num_mamba_layers > 0: - first_mamba_local_layer = self.mamba_local_layer_ids[0] + first_mamba_local_layer = self.layer_offsets[ + self.mamba_pp_layers[0]] self.ssm_layer_group_id = self.impl.get_layer_group_id( LayerId(first_mamba_local_layer)) self._ssm_page_index_scale = self.impl.get_page_index_scale( @@ -2736,7 +2722,6 @@ def get_cache_size_per_token(model_config, mapping, max_batch_size=max_batch_size, kv_cache_config=kv_cache_config, - num_layers=num_layers, tokens_per_block=tokens_per_block, max_seq_len=max_seq_len, num_reserved_dummy_slots=num_reserved_dummy_slots, @@ -2838,10 +2823,6 @@ def _planned_token_capacity( gpu_quota: int, ) -> int: capacity = self._get_max_tokens_from_quota(gpu_quota) - if math.isinf(capacity): - capacity = (kv_cache_config.max_tokens if kv_cache_config.max_tokens - is not None else self.max_seq_len * - self._max_resident_sequences()) capacity = min( capacity, self.max_seq_len * self._max_resident_sequences(), @@ -2852,26 +2833,19 @@ def _planned_token_capacity( def _minimum_live_gpu_quota(self) -> int: """Return the minimum quota for live states and one attention page.""" - num_sequences = self._max_resident_sequences() - state_slots = num_sequences + self._num_reserved_dummy_slots - state_quota = state_slots * self._mamba_state_bytes_per_slot() attention_block_quota = (self._attention_cache_bytes_per_token() * self.tokens_per_block) - # At zero token capacity, any reserved dummy state makes non-live SSM - # capacity possible. Match the normal estimator by reserving one - # partial attention page per request lineage in that case. - extra_attention_quota = (num_sequences * attention_block_quota - if state_slots > num_sequences else 0) - attention_quota = max( - super()._get_quota_from_max_tokens(0) + extra_attention_quota, - attention_block_quota, + num_state_slots = (self._max_resident_sequences() + + self._num_reserved_dummy_slots) + state_quota = num_state_slots * self._mamba_state_bytes_per_slot() + return max( + self._get_quota_from_max_tokens(0), + state_quota + attention_block_quota, ) - return state_quota + attention_quota def _build_cache_config( self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: kv_cache_config = self.kv_cache_config - tokens_per_block = config.tokens_per_block cache_tiers = config.cache_tiers gpu_quota = cache_tiers[0].quota minimum_live_quota = self._minimum_live_gpu_quota() @@ -2880,51 +2854,22 @@ def _build_cache_config( "The V2 Mamba GPU cache quota is too small for live recurrent " f"states and attention pages: got {gpu_quota} bytes, need at " f"least {minimum_live_quota} bytes.") - buffer_type = [Role.KEY] - if self.kv_cache_type != CacheTypeCpp.SELFKONLY: - buffer_type.append(Role.VALUE) - if kv_cache_config.dtype == "nvfp4": - for layer_idx, hd in enumerate(self.head_dim_per_layer): - if self._is_local_mamba_layer(layer_idx): - continue - assert hd % 2 == 0, ( - f"head_dim must be divisible by 2 for nvfp4 kv cache, " - f"but layer {layer_idx} has head_dim={hd}") - buffer_type.append(Role.KEY_BLOCK_SCALE) - if self.kv_cache_type != CacheTypeCpp.SELFKONLY: - buffer_type.append(Role.VALUE_BLOCK_SCALE) - - layers = [] + # _build_base_config already constructed every attention layer, + # including dtype-specific scale and subclass-provided side buffers. + # Preserve those configs and replace only the local Mamba layers. + layers = list(config.layers) for local_layer_idx, global_layer_idx in enumerate(self.pp_layers): - layer_id = LayerId(local_layer_idx) if self._mamba_layer_mask[global_layer_idx]: - layers.append( - SsmLayerConfig( - layer_id=layer_id, - buffers=[ - BufferConfig(role=MambaRole.SSM_STATE, - size=self.ssm_bytes), - BufferConfig(role=MambaRole.CONV_STATE, - size=self.conv_bytes), - ], - )) - else: - layers.append( - AttentionLayerConfig( - layer_id=layer_id, - buffers=[ - BufferConfig( - role=role, - size=self.get_layer_bytes_per_token( - local_layer_idx=layer_id, data_role=role) * - tokens_per_block, - ) for role in buffer_type - ], - sliding_window_size=self.max_attention_window_vec[ - global_layer_idx % - len(self.max_attention_window_vec)], - num_sink_tokens=None, - )) + layer_id = LayerId(local_layer_idx) + layers[local_layer_idx] = SsmLayerConfig( + layer_id=layer_id, + buffers=[ + BufferConfig(role=MambaRole.SSM_STATE, + size=self.ssm_bytes), + BufferConfig(role=MambaRole.CONV_STATE, + size=self.conv_bytes), + ], + ) num_sequences = self._max_resident_sequences() planned_capacity = self._planned_token_capacity(kv_cache_config, @@ -2972,15 +2917,18 @@ def _get_state_buffer(self, local_layer_idx: int, role, dtype: torch.dtype, ) def _setup_states(self) -> None: + local_layer_ids = [ + self.layer_offsets[layer_id] for layer_id in self.mamba_pp_layers + ] self.all_ssm_states = [ self._get_state_buffer(local_layer_idx, MambaRole.SSM_STATE, self.ssm_state_dtype, self.ssm_state_shape) - for local_layer_idx in self.mamba_local_layer_ids + for local_layer_idx in local_layer_ids ] self.all_conv_states = [ self._get_state_buffer(local_layer_idx, MambaRole.CONV_STATE, self.conv_state_dtype, self.conv_state_shape) - for local_layer_idx in self.mamba_local_layer_ids + for local_layer_idx in local_layer_ids ] def _setup_replay_buffers(self, spec_config) -> None: @@ -3016,13 +2964,9 @@ def get_cache_bytes_per_token(self) -> int: self.kv_cache_config.mamba_state_config.periodic_snapshot_interval) if (self.kv_cache_config.enable_block_reuse and interval is not None and interval > 0): - cache_bytes += (self.local_num_mamba_layers * - (self.ssm_bytes + self.conv_bytes) // interval) + cache_bytes += self._mamba_state_bytes_per_slot() // interval if cache_bytes == 0 and self.local_num_mamba_layers > 0: - return max( - 1, - self.local_num_mamba_layers * - (self.ssm_bytes + self.conv_bytes)) + cache_bytes = self._mamba_state_bytes_per_slot() return max(1, cache_bytes) def get_num_free_blocks(self) -> int: diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py b/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py index 0a05f393346f..a37288a80895 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py @@ -120,8 +120,13 @@ def test_capacity_only_is_scoped_to_target_manager() -> None: target_cache.resize.assert_called_once_with(253, None) -def test_dynamic_tree_draft_reclaims_reserved_capacity() -> None: - manager = _manager(is_draft=True, kv_reserve_draft_tokens=60) +@pytest.mark.parametrize( + ("is_draft", "expected_capacity"), + [(True, 201), (False, 230)], + ids=["draft-reclaims-reserve", "target-has-no-reserve"], +) +def test_dynamic_tree_reserved_capacity(is_draft: bool, expected_capacity: int) -> None: + manager = _manager(is_draft=is_draft, kv_reserve_draft_tokens=60) # The runtime tree used 31 draft positions: 26 rejected and 5 accepted. request = _request(1, rewind=26, accepted_draft_tokens=5) cache = _cache() @@ -129,19 +134,7 @@ def test_dynamic_tree_draft_reclaims_reserved_capacity() -> None: manager.update_resources(SimpleNamespace(generation_requests=[request])) - # Rewind 26 rejected positions plus the 60 - 31 reserve slack. - cache.resize.assert_called_once_with(201, 200) - - -def test_dynamic_tree_target_does_not_reclaim_unallocated_reserve() -> None: - manager = _manager(is_draft=False, kv_reserve_draft_tokens=60) - request = _request(1, rewind=26, accepted_draft_tokens=5) - cache = _cache() - manager.kv_cache_map[request.py_request_id] = cache - - manager.update_resources(SimpleNamespace(generation_requests=[request])) - - cache.resize.assert_called_once_with(230, 200) + cache.resize.assert_called_once_with(expected_capacity, 200) def test_capacity_only_completion_preserves_history() -> None: diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 1d0dbe3ad7d7..451cace348c1 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -66,11 +66,14 @@ ) from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + AttentionLayerConfig, BatchDesc, + BufferConfig, GpuCacheTierConfig, KVCacheDesc, KVCacheManagerConfig, LayerId, + SsmLayerConfig, ) from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManager as RuntimeKVCacheManager @@ -194,9 +197,7 @@ def test_mamba_kv_cache_params_separate_target_and_draft_masks(): [ (True, False, MambaHybridCacheManagerV2), (True, True, MambaHybridCacheManagerV2), - (False, False, CppMambaHybridCacheManager), (False, True, CppMambaHybridCacheManager), - ("auto", False, CppMambaHybridCacheManager), ("auto", True, CppMambaHybridCacheManager), ], ) @@ -362,64 +363,6 @@ def test_hybrid_cache_manager_factory_requires_v2_for_explicit_snapshots( ) -@pytest.mark.parametrize( - ("use_v2", "expected"), - [ - (False, CppMambaHybridCacheManager), - ("auto", CppMambaHybridCacheManager), - (True, MambaHybridCacheManagerV2), - ], -) -def test_hybrid_cache_manager_factory_selects_manager_when_snapshot_reuse_disabled( - monkeypatch, use_v2, expected -): - monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) - monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) - - assert ( - get_kv_cache_manager_cls( - _hybrid_model_config(), - KvCacheConfig( - enable_block_reuse=False, - mamba_state_config=MambaStateConfig(periodic_snapshot_interval=0), - use_kv_cache_manager_v2=use_v2, - ), - ) - is expected - ) - - -@pytest.mark.parametrize( - ("field", "offsets"), - [ - ("additional_snapshot_offsets_from_start", [128]), - ("additional_snapshot_offsets_from_end", [0]), - ], -) -def test_hybrid_cache_manager_factory_routes_explicit_snapshots_to_v2( - monkeypatch, - field, - offsets, -): - monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) - monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) - - assert ( - get_kv_cache_manager_cls( - _hybrid_model_config(), - KvCacheConfig( - enable_block_reuse=True, - mamba_state_config=MambaStateConfig( - periodic_snapshot_interval=0, - **{field: offsets}, - ), - use_kv_cache_manager_v2=True, - ), - ) - is MambaHybridCacheManagerV2 - ) - - @pytest.mark.parametrize("backend", ["NIXL", "DEFAULT"]) def test_hybrid_cache_manager_factory_routes_explicit_v2_disagg(monkeypatch, backend): monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) @@ -1242,32 +1185,6 @@ def test_v2_hybrid_estimator_accounts_for_ssm_slot_attention_bound( ) == (11, expected_intercept) -def test_v2_hybrid_estimator_reserves_attention_for_each_lineage(monkeypatch): - monkeypatch.setattr( - KVCacheManager, - "get_cache_size_per_token", - lambda *args, **kwargs: 11, - ) - monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.mamba_cache_manager._get_local_mamba_cache_layout", - lambda *args, **kwargs: ( - SimpleNamespace(get_states_bytes_per_layer=lambda mapping: 64), - 1, - 1, - ), - ) - - # The one CUDA-graph padding slot is less than the four resident request - # lineages, but the conservative attention bound still reserves one page - # for every lineage. - assert MambaHybridCacheManagerV2.get_cache_size_per_token( - object(), - Mapping(world_size=1, tp_size=1, pp_size=1), - max_batch_size=4, - kv_cache_config=KvCacheConfig(enable_block_reuse=False), - ) == (11, 1728) - - def test_v2_hybrid_attention_bound_is_snapshot_alignment_agnostic(): model_config = _hybrid_cache_sizing_model_config(["linear_attention", "full_attention"]) mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) @@ -1303,6 +1220,16 @@ def test_v2_hybrid_planned_capacity_is_bounded_by_resident_sequences(): assert capacity == 4 * 2 * 128 +def _base_attention_layer_configs(num_layers): + return [ + AttentionLayerConfig( + layer_id=LayerId(layer_idx), + buffers=[BufferConfig(role="key", size=256)], + ) + for layer_idx in range(num_layers) + ] + + def test_v2_hybrid_typical_batch_uses_num_ssm_slots(): mgr = object.__new__(MambaHybridCacheManagerV2) mgr.kv_cache_type = CacheTypeCpp.SELF @@ -1333,15 +1260,18 @@ def test_v2_hybrid_typical_batch_uses_num_ssm_slots(): ) mgr.kv_cache_config = kv_cache_config constraints = [BatchDesc([KVCacheDesc(capacity=64, history_length=0)])] + base_layers = _base_attention_layer_configs(2) base_config = KVCacheManagerConfig( tokens_per_block=32, cache_tiers=[GpuCacheTierConfig(quota=1 << 20)], - layers=[], + layers=base_layers, constraints=constraints, ) config = mgr._build_cache_config(base_config) + assert isinstance(config.layers[0], SsmLayerConfig) + assert config.layers[1] is base_layers[1] assert len(config.typical_step.kv_caches) == 2 assert [kv.num_ssm_slots for kv in config.typical_step.kv_caches] == [3, 2] assert config.constraints is constraints @@ -1366,9 +1296,9 @@ def test_v2_hybrid_rejects_quota_below_live_state_floor(): mgr.enable_swa_scratch_reuse = False mgr.get_layer_bytes_per_token = lambda **kwargs: 8 mgr._attention_cache_bytes_per_token = lambda: 16 + mgr.kv_cache_config = KvCacheConfig(enable_partial_reuse=False) minimum_quota = mgr._minimum_live_gpu_quota() - mgr.kv_cache_config = KvCacheConfig(enable_partial_reuse=False) base_config = KVCacheManagerConfig( tokens_per_block=32, cache_tiers=[GpuCacheTierConfig(quota=minimum_quota - 1)], @@ -1395,6 +1325,7 @@ def test_v2_hybrid_pure_mamba_rank_does_not_reserve_attention_page(): mgr.enable_swa_scratch_reuse = False mgr.get_layer_bytes_per_token = lambda **kwargs: 0 mgr._attention_cache_bytes_per_token = lambda: 0 + mgr.kv_cache_config = KvCacheConfig(enable_block_reuse=False) assert mgr._minimum_live_gpu_quota() == 3 * (64 + 32) @@ -1526,31 +1457,19 @@ def test_v2_hybrid_add_dummy_requests_forwards_encoder_output_lens(mocker): assert base_add_dummy_requests.call_args.kwargs["encoder_output_lens"] == [17] -def test_v2_hybrid_prepare_expect_snapshot_points_clears_when_reuse_disabled(): - mgr = object.__new__(MambaHybridCacheManagerV2) - mgr.enable_block_reuse = False - mgr.kv_cache_config = KvCacheConfig(enable_block_reuse=False) - request = SimpleNamespace(prompt_len=150, expect_snapshot_points=[64]) - - mgr.prepare_expect_snapshot_points([request]) - - assert request.expect_snapshot_points == [] - - -def test_cpp_hybrid_prepare_expect_snapshot_points_clears_when_reuse_disabled(): - mgr = object.__new__(CppMambaHybridCacheManager) +@pytest.mark.parametrize( + "manager_cls", + [ + MambaHybridCacheManagerV2, + CppMambaHybridCacheManager, + MixedMambaHybridCacheManager, + ], + ids=["v2", "cpp", "mixed"], +) +def test_hybrid_prepare_expect_snapshot_points_clears_when_reuse_disabled(manager_cls): + mgr = object.__new__(manager_cls) mgr.enable_block_reuse = False mgr.kv_cache_config = KvCacheConfig(enable_block_reuse=False) - request = SimpleNamespace(prompt_len=150, expect_snapshot_points=[64]) - - mgr.prepare_expect_snapshot_points([request]) - - assert request.expect_snapshot_points == [] - - -def test_mixed_hybrid_snapshot_hook_clears_when_reuse_disabled(): - mgr = object.__new__(MixedMambaHybridCacheManager) - mgr.enable_block_reuse = False request = SimpleNamespace(prompt_len=64, expect_snapshot_points=[64]) mgr.prepare_expect_snapshot_points([request]) @@ -1558,13 +1477,12 @@ def test_mixed_hybrid_snapshot_hook_clears_when_reuse_disabled(): assert request.expect_snapshot_points == [] -@pytest.mark.parametrize("interval", [0, -64, None]) -def test_cpp_hybrid_prepare_expect_snapshot_points_clears_for_invalid_interval(interval): +def test_cpp_hybrid_prepare_expect_snapshot_points_clears_for_disabled_interval(): mgr = object.__new__(CppMambaHybridCacheManager) mgr.enable_block_reuse = True mgr.kv_cache_config = SimpleNamespace( enable_block_reuse=True, - mamba_state_config=SimpleNamespace(periodic_snapshot_interval=interval), + mamba_state_config=SimpleNamespace(periodic_snapshot_interval=0), ) request = SimpleNamespace(prompt_len=150, expect_snapshot_points=[64]) @@ -1621,7 +1539,7 @@ def allocated_memory(pool_ratio): base_config = KVCacheManagerConfig( tokens_per_block=32, cache_tiers=[GpuCacheTierConfig(quota=64 << 20)], - layers=[], + layers=_base_attention_layer_configs(2), initial_pool_ratio=pool_ratio, ) config = mgr._build_cache_config(base_config) @@ -1929,8 +1847,8 @@ def test_v2_hybrid_invalid_check_scans_distinct_attention_pools(): @skip_no_cuda -@pytest.mark.parametrize("max_batch_size", [1, 4]) -def test_v2_hybrid_dummy_indices_keep_cuda_buffer_address(max_batch_size): +def test_v2_hybrid_dummy_indices_keep_cuda_buffer_address(): + max_batch_size = 1 mgr = _build_v2_hybrid_with_mamba_layer(max_batch_size=max_batch_size, enable_attention_dp=True) try: request_ids = list(range(100, 100 + max_batch_size)) @@ -1976,8 +1894,6 @@ def test_v2_hybrid_reserves_every_persistent_dummy_slot(): ] request_ids = [101, 102, 103, 104] - assert mgr._num_cuda_graph_padding_dummy_slots == 4 - assert mgr._num_attention_dp_dummy_slots == 1 assert mgr._num_reserved_dummy_slots == 5 assert mgr.index_mapper.num_free_slots() == len(request_ids) + 5 @@ -2123,14 +2039,15 @@ def test_v2_hybrid_mamba_state_views_use_logical_slots(): assert all(t.shape[0] == conv_slots for t in mgr.all_conv_states) assert ssm_slots == conv_slots + local_layer_ids = [mgr.layer_offsets[layer_id] for layer_id in mgr.mamba_pp_layers] for local_layer_idx, ssm_state, conv_state in zip( - mgr.mamba_local_layer_ids, mgr.all_ssm_states, mgr.all_conv_states + local_layer_ids, mgr.all_ssm_states, mgr.all_conv_states ): layer_id = LayerId(local_layer_idx) ssm_scale = mgr.impl.get_page_index_scale(layer_id, MambaRole.SSM_STATE) conv_scale = mgr.impl.get_page_index_scale(layer_id, MambaRole.CONV_STATE) - assert ssm_state.stride(0) == mgr.ssm_count * ssm_scale - assert conv_state.stride(0) == mgr.conv_count * conv_scale + assert ssm_state.stride(0) == ssm_state[0].numel() * ssm_scale + assert conv_state.stride(0) == conv_state[0].numel() * conv_scale assert ( ssm_state.shape[0] == ( @@ -2232,30 +2149,6 @@ def test_v2_hybrid_disagg_page_table_uses_qwen3_next_conv_sections(): mgr.shutdown() -@skip_no_cuda -def test_cpp_hybrid_replay_buffers_size_by_tokens_per_gen_step(): - spec_config = _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5) - mgr = _build_hybrid_with_mamba_layer( - spec_config=spec_config, - max_batch_size=4, - use_replay_state_update=True, - ) - try: - replay_metadata = mgr.get_replay_state_update_metadata() - assert mgr.use_replay_state_update is True - assert replay_metadata is not None - assert replay_metadata.replay_step_width == spec_config.tokens_per_gen_step - assert replay_metadata.replay_history_size == max( - MIN_REPLAY_HISTORY_SIZE, spec_config.tokens_per_gen_step - ) - layer_cache = mgr.mamba_layer_cache(0) - _assert_replay_layer_cache_uses_history_size( - layer_cache, replay_metadata.replay_history_size - ) - finally: - mgr.shutdown() - - @skip_no_cuda def test_v2_hybrid_intermediate_states_size_by_tokens_per_gen_step(): spec_config = _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5) @@ -2290,16 +2183,21 @@ def test_v2_hybrid_static_dynamic_tree_capacity(): assert mgr.intermediate_ssm_states.shape[2] == 32 assert mgr.intermediate_conv_states.shape[2] == 32 assert mgr._kv_reserve_draft_tokens == 31 - assert mgr._num_cuda_graph_padding_dummy_slots == 1 + assert mgr._num_reserved_dummy_slots == 1 assert not mgr.use_replay_state_update finally: mgr.shutdown() @skip_no_cuda -def test_v2_hybrid_replay_buffers_size_by_tokens_per_gen_step(): +@pytest.mark.parametrize( + "builder", + [_build_hybrid_with_mamba_layer, _build_v2_hybrid_with_mamba_layer], + ids=["cpp", "v2"], +) +def test_hybrid_replay_buffers_size_by_tokens_per_gen_step(builder): spec_config = _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5) - mgr = _build_v2_hybrid_with_mamba_layer( + mgr = builder( max_batch_size=4, spec_config=spec_config, use_replay_state_update=True, diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index 7fca29ef2f81..69fb84b26887 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -610,10 +610,6 @@ def test_mamba_layer_group_serialization(): d = mlg.to_dict() assert d["mamba_layer_offsets"] == {10: 0, 11: 1, 12: 2} assert d["conv_section_bytes"] == [512, 256, 256] - assert d["conv_states"]["slot_stride_bytes"] == 512 - assert d["conv_states"]["layer_stride_bytes"] == 256 - assert d["ssm_states"]["slot_stride_bytes"] == 1024 - assert d["ssm_states"]["layer_stride_bytes"] == 512 from tensorrt_llm._torch.disaggregation.resource.page import LayerGroup @@ -621,16 +617,10 @@ def test_mamba_layer_group_serialization(): assert isinstance(restored, MambaLayerGroup) assert restored.mamba_layer_offsets == {10: 0, 11: 1, 12: 2} assert restored.conv_states.base_address == 1000 - assert restored.conv_states.slot_bytes == 128 - assert restored.conv_states.num_slots == 10 assert restored.conv_states.slot_stride_bytes == 512 - assert restored.conv_states.layer_stride_bytes == 256 assert get_slot_address(restored.conv_states, 3) == 1000 + 3 * 512 assert restored.ssm_states.base_address == 8000 - assert restored.ssm_states.slot_bytes == 256 - assert restored.ssm_states.num_slots == 8 assert restored.ssm_states.slot_stride_bytes == 1024 - assert restored.ssm_states.layer_stride_bytes == 512 assert restored.conv_section_bytes == [512, 256, 256] assert restored.ssm_bytes_per_head == 128 diff --git a/tests/unittest/disaggregated/test_mamba_transfer.py b/tests/unittest/disaggregated/test_mamba_transfer.py index c9aee4d5a3d3..790ce1aed528 100644 --- a/tests/unittest/disaggregated/test_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_mamba_transfer.py @@ -24,7 +24,6 @@ import tensorrt_llm.bindings import tensorrt_llm.tensorrt_llm_transfer_agent_binding # noqa: F401 from tensorrt_llm import DisaggregatedParams, Mapping, SamplingParams -from tensorrt_llm._torch.disaggregation.native import rank_info from tensorrt_llm._torch.disaggregation.native.mixers.ssm import peer from tensorrt_llm._torch.disaggregation.resource import page from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 @@ -256,7 +255,7 @@ def _mamba_layer_ids(manager): def _mamba_state_slot(manager, request_id): if isinstance(manager, MambaHybridCacheManagerV2): - return manager.get_state_indices([request_id])[0] + return manager._request_id_to_state_index[request_id] return manager.mamba_cache_index[request_id] @@ -279,89 +278,6 @@ def test_mamba_policy_layer_major_v1_ptrs(): np.testing.assert_array_equal(ptrs, [130, 210]) -def test_mamba_policy_slot_major_layer_ptrs(): - """V2 Mamba state tensors step by physical slots, not V1 layers.""" - self_mlg = page.MambaLayerGroup( - pool_group_idx=0, - mamba_layer_offsets={1: 0, 2: 1}, - conv_states=page.PhysicalPool( - base_address=1000, - slot_bytes=10, - num_slots=8, - slot_stride_bytes=20, - layer_stride_bytes=10, - ), - ssm_states=page.PhysicalPool( - base_address=3000, - slot_bytes=20, - num_slots=8, - slot_stride_bytes=40, - layer_stride_bytes=20, - ), - conv_section_bytes=[10], - ssm_bytes_per_head=10, - ) - peer_mlg = page.MambaLayerGroup( - pool_group_idx=0, - mamba_layer_offsets={1: 0, 2: 1}, - conv_states=page.PhysicalPool( - base_address=5000, - slot_bytes=10, - num_slots=8, - slot_stride_bytes=20, - layer_stride_bytes=10, - ), - ssm_states=page.PhysicalPool( - base_address=7000, - slot_bytes=20, - num_slots=8, - slot_stride_bytes=40, - layer_stride_bytes=20, - ), - conv_section_bytes=[10], - ssm_bytes_per_head=10, - ) - self_ri = rank_info.RankInfo( - instance_name="self", - instance_rank=0, - tp_size=1, - tp_rank=0, - pp_size=1, - pp_rank=0, - layer_num_per_pp=[2], - sender_endpoints=[], - server_endpoint="", - self_endpoint="", - transfer_engine_info=bytes(), - ) - peer_ri = rank_info.RankInfo( - instance_name="peer", - instance_rank=0, - tp_size=1, - tp_rank=0, - pp_size=1, - pp_rank=0, - layer_num_per_pp=[2], - sender_endpoints=[], - server_endpoint="", - self_endpoint="", - transfer_engine_info=bytes(), - ) - - src_frags, dst_frags, kv_sizes = peer.MambaPolicy.build_mamba_frags( - self_mlg=self_mlg, - peer_mlg=peer_mlg, - src_slot=3, - dst_slot=5, - self_ri=self_ri, - peer_ri=peer_ri, - ) - - assert src_frags == [1060, 1070, 3120, 3140] - assert dst_frags == [5100, 5110, 7200, 7220] - assert kv_sizes == [10, 10, 20, 20] - - def test_mamba_policy_slot_major_interleaved_role_ptrs(): """Layer/role offsets remain inside each coalesced V2 physical slot.""" state_bytes = 64 diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index 7dba91291797..36a49e73f8ec 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -2770,7 +2770,6 @@ def test_num_ssm_slots_controls_ssm_slot_count(self): ssm_lc = manager._life_cycles.ssm_life_cycle_id assert ssm_lc is not None ssm_pg = manager._storage.get_pool_group_index(ssm_lc) - attn_pg = 1 - ssm_pg batch = BatchDesc( kv_caches=[ @@ -2784,10 +2783,6 @@ def test_num_ssm_slots_controls_ssm_slot_count(self): ) slots = manager._storage._compute_slots_for_batch(batch, self.TOKENS_PER_BLOCK, None) self.assertEqual(slots[ssm_pg], 6) - self.assertEqual( - slots[attn_pg], - 2 * div_up(1024, self.TOKENS_PER_BLOCK) + 2, - ) default_batch = BatchDesc(kv_caches=[KVCacheDesc(capacity=4096, history_length=4095)] * 2) default_slots = manager._storage._compute_slots_for_batch( diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 985944a96ed3..2478b50f5ee3 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -700,17 +700,6 @@ def test_KvCacheConfig_declaration(): assert pybind_config.attention_dp_events_gather_period_ms == 10 assert (KvCacheConfig(block_reuse_policy="per_conversation"). block_reuse_policy == "per_conversation") - per_conversation_config = KvCacheConfig( - block_reuse_policy="per_conversation", - mamba_state_config=MambaStateConfig( - periodic_snapshot_interval=64, - additional_snapshot_offsets_from_end=[0], - ), - ) - assert (per_conversation_config.mamba_state_config. - periodic_snapshot_interval == 0) - assert (per_conversation_config.mamba_state_config. - additional_snapshot_offsets_from_end == [0]) with pytest.raises(ValidationError): KvCacheConfig(block_reuse_policy="invalid") @@ -738,14 +727,9 @@ def test_MambaStateConfig_rejects_unknown_fields(): [ ("periodic_snapshot_interval", -1), ("additional_snapshot_offsets_from_start", [0]), - ("additional_snapshot_offsets_from_start", [-1]), ("additional_snapshot_offsets_from_start", [True]), - ("additional_snapshot_offsets_from_start", [1.5]), - ("additional_snapshot_offsets_from_start", ["128"]), ("additional_snapshot_offsets_from_end", [-1]), - ("additional_snapshot_offsets_from_end", [True]), ("additional_snapshot_offsets_from_end", [1.5]), - ("additional_snapshot_offsets_from_end", ["32"]), ], ) def test_MambaStateConfig_rejects_invalid_snapshot_offsets(field, value): @@ -770,12 +754,11 @@ def test_KvCacheConfig_requires_v2_for_additional_snapshot_offsets( use_kv_cache_manager_v2=False, ) - for use_v2 in (True, "auto"): - config = KvCacheConfig( - mamba_state_config=state_config, - use_kv_cache_manager_v2=use_v2, - ) - assert getattr(config.mamba_state_config, field) == offsets + config = KvCacheConfig( + mamba_state_config=state_config, + use_kv_cache_manager_v2=True, + ) + assert getattr(config.mamba_state_config, field) == offsets def test_KvCacheConfig_allows_periodic_snapshots_with_v1(): @@ -800,12 +783,6 @@ def test_KvCacheConfig_migrates_deprecated_mamba_interval(monkeypatch): for message in warnings_seen) assert "mamba_state_cache_interval" not in config.model_dump() - llm_args = TorchLlmArgs( - model=llama_model_path, - kv_cache_config={"mamba_state_cache_interval": 32}, - ) - assert llm_args.kv_cache_config.mamba_state_config.periodic_snapshot_interval == 32 - def test_KvCacheConfig_warns_when_disabling_periodic_conversation_snapshots( monkeypatch): @@ -815,10 +792,14 @@ def test_KvCacheConfig_warns_when_disabling_periodic_conversation_snapshots( config = KvCacheConfig( block_reuse_policy="per_conversation", - mamba_state_config=MambaStateConfig(periodic_snapshot_interval=64), + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=64, + additional_snapshot_offsets_from_end=[0], + ), ) assert config.mamba_state_config.periodic_snapshot_interval == 0 + assert config.mamba_state_config.additional_snapshot_offsets_from_end == [0] assert len(warnings_seen) == 1 assert "periodic_snapshot_interval=64" in warnings_seen[0] assert "block_reuse_policy=per_conversation" in warnings_seen[0] From 8f6c51fec27c23cc0f8620df205fecc1c0711aaa Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Fri, 24 Jul 2026 08:29:52 +0800 Subject: [PATCH 21/22] [Agent fix] Fix stale V2 Mamba replay checkpoint test (By Agent) test_v2_hybrid_replay_bookkeeping_matches_checkpoint_predicate hardcoded a checkpoint-crossing scenario assuming replay_history_size == tokens_per_gen_step (5). Since replay_history_size became max(MIN_REPLAY_HISTORY_SIZE=16, tokens_per_gen_step), the scenario no longer crossed the checkpoint threshold, so prev_num_accepted_tokens accumulated (3+2=5) instead of resetting to 2. Derive the threshold from the manager's replay metadata and drive update_mamba_states until it actually crosses prev + step_width > history_size, so the test tracks the real checkpoint predicate regardless of the minimum history size. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../executor/test_mamba_cache_manager.py | 41 +++++++++++++------ 1 file changed, 29 insertions(+), 12 deletions(-) diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 451cace348c1..b2728622667e 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -2231,22 +2231,39 @@ def test_v2_hybrid_replay_bookkeeping_matches_checkpoint_predicate(monkeypatch): lambda *args, **kwargs: None, ) try: + # Derive the checkpoint threshold from the manager's own replay + # metadata instead of hardcoding it: replay_history_size is + # max(MIN_REPLAY_HISTORY_SIZE, tokens_per_gen_step), so the boundary + # is not simply tokens_per_gen_step. + replay_metadata = mgr.get_replay_state_update_metadata() + step_width = replay_metadata.replay_step_width + history_size = replay_metadata.replay_history_size + slot = torch.tensor([0], dtype=torch.int32, device="cuda") attn_metadata = SimpleNamespace(num_seqs=1, num_contexts=0) - mgr.update_mamba_states( - attn_metadata, - torch.tensor([3], dtype=torch.int32, device="cuda"), - state_indices=slot, - ) - assert mgr.prev_num_accepted_tokens[0].item() == 3 - assert mgr.cache_buf_idx[0].item() == 0 + def advance(accepted): + mgr.update_mamba_states( + attn_metadata, + torch.tensor([accepted], dtype=torch.int32, device="cuda"), + state_indices=slot, + ) - mgr.update_mamba_states( - attn_metadata, - torch.tensor([2], dtype=torch.int32, device="cuda"), - state_indices=slot, - ) + # Accumulate below the checkpoint threshold: while + # prev + step_width <= history_size the manager keeps writing into the + # same cache_buf_idx and grows prev_num_accepted_tokens monotonically. + prev = 0 + while prev + step_width <= history_size: + advance(1) + prev += 1 + assert mgr.prev_num_accepted_tokens[0].item() == prev + assert mgr.cache_buf_idx[0].item() == 0 + + # The next step crosses the threshold (prev + step_width > history_size): + # the manager starts a fresh checkpoint, resetting + # prev_num_accepted_tokens to the accepted count and flipping + # cache_buf_idx to the other buffer. + advance(2) assert mgr.prev_num_accepted_tokens[0].item() == 2 assert mgr.cache_buf_idx[0].item() == 1 finally: From 50520cb7355bc08e9a947fd988e571f4b48914fd Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Sat, 25 Jul 2026 11:29:54 +0800 Subject: [PATCH 22/22] [None][fix] align V2 Mamba cache sizing (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- docs/source/features/kvcache.md | 7 +- .../_torch/pyexecutor/mamba_cache_manager.py | 123 +++++++++------ .../_torch/pyexecutor/model_loader.py | 21 +-- tensorrt_llm/llmapi/llm_args.py | 16 +- tensorrt_llm/llmapi/llm_utils.py | 27 +++- .../runtime/kv_cache_manager_v2/__init__.pyi | 1 - .../runtime/kv_cache_manager_v2/_config.py | 2 - .../kv_cache_manager_v2/_core/_kv_cache.py | 14 +- .../kv_cache_manager_v2/_storage_manager.py | 19 +-- .../defs/accuracy/test_llm_api_pytorch.py | 11 ++ .../executor/test_mamba_cache_manager.py | 143 +++++++----------- .../test_kv_cache_manager_v2.py | 85 +---------- tests/unittest/llmapi/test_llm_args.py | 79 ++++++++-- 13 files changed, 269 insertions(+), 279 deletions(-) diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index f8b8afc14ddf..82f7fbe58143 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -96,6 +96,7 @@ an end offset of `0`) when using it with a hybrid Mamba model. For example: kv_cache_config: enable_block_reuse: true use_kv_cache_manager_v2: true + avg_seq_len: 2048 mamba_state_config: periodic_snapshot_interval: 0 additional_snapshot_offsets_from_start: [128] @@ -104,7 +105,11 @@ kv_cache_config: This retains snapshots after the first 128 tokens, at the end of the prompt, and before the final 32 prompt tokens. Positions outside a particular prompt -are ignored. Exact explicit boundaries currently require +are ignored. Set `avg_seq_len` to the workload's average total sequence length +so V2 can size the attention KV and Mamba state pools in the right proportion. +If neither `avg_seq_len` nor an explicit `pool_ratio` is configured, hybrid +Mamba models warn and fall back to half of `max_seq_len`, which can produce a +suboptimal pool split. Exact explicit boundaries currently require `MambaHybridCacheManagerV2`, `max_beam_width=1`, and no KV connector. Hybrid Mamba models select V2 by default when `use_kv_cache_manager_v2: auto`; set it to `false` to select the V1 C++ diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 425d26fc969d..e7cab986c2ef 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -1699,13 +1699,10 @@ def _estimate_mamba_hybrid_cache_cost( has_unaligned_periodic_snapshot = (interval is not None and interval % tokens_per_block != 0) if cap_partial_attention_snapshots: - # V2 storage receives only an SSM slot count, which does not reveal - # whether each snapshot is block-aligned. Once any non-live SSM slot is - # possible, reserve one retained partial attention page per resident - # lineage. Fewer non-live slots and dummy slots only make this an - # overestimate; the bound does not rely on S >= 2N. - has_non_live_ssm_capacity = (fixed_rules > 0 or interval is not None - or num_reserved_dummy_slots > 0) + # Snapshot alignment is unknown while estimating cache cost. Once a + # snapshot is possible, reserve one retained partial attention page + # per resident lineage. Dummy requests carry no attention capacity. + has_non_live_ssm_capacity = fixed_rules > 0 or interval is not None partial_attention_slots = (max_resident_sequences if has_non_live_ssm_capacity else 0) else: @@ -2763,30 +2760,65 @@ def _num_ssm_snapshots_for_capacity( return (self._max_resident_sequences() * fixed_rules + regular_snapshots) - def _ssm_slots_per_request_for_typical_batch( + def _num_ssm_states_per_typical_request( self, capacity: int, kv_cache_config: KvCacheConfig, - ) -> List[int]: - num_sequences = self._max_resident_sequences() - snapshot_slots = self._num_ssm_snapshots_for_capacity( + ) -> int: + fixed_rules, _ = _mamba_snapshot_rule_counts( + kv_cache_config, + capacity, + self.tokens_per_block, + ) + # Additional snapshots are stable boundaries that must remain alive. + # Periodic snapshots are evictable cache entries and therefore do not + # increase the guaranteed state count represented by BatchDesc. + return 1 + fixed_rules + + def _typical_request_descs( + self, + capacity: int, + kv_cache_config: KvCacheConfig, + ) -> List[KVCacheDesc]: + """Model one request with one descriptor per live SSM state.""" + num_states = self._num_ssm_states_per_typical_request( capacity, kv_cache_config) - snapshot_slots_per_request, snapshot_remainder = divmod( - snapshot_slots, num_sequences) - dummy_per_request, dummy_remainder = divmod( - self._num_reserved_dummy_slots, num_sequences) + capacity_per_state, capacity_remainder = divmod(capacity, num_states) + capacities = [ + capacity_per_state + int(i < capacity_remainder) + for i in range(num_states) + ] return [ - 1 + snapshot_slots_per_request + int(i < snapshot_remainder) + - dummy_per_request + int(i < dummy_remainder) - for i in range(num_sequences) + KVCacheDesc( + capacity=state_capacity, + history_length=max(0, state_capacity - 1), + ) for state_capacity in capacities ] + def _get_typical_request_capacity( + self, + kv_cache_config: KvCacheConfig, + ) -> int: + if kv_cache_config.avg_seq_len is not None: + return kv_cache_config.avg_seq_len + + fallback_capacity = max(1, self.max_seq_len // 2) + logger.warning( + "'kv_cache_config.avg_seq_len' is not set for a hybrid Mamba " + "model using KV cache manager V2. Falling back to " + f"max_seq_len / 2={fallback_capacity} for cache-pool sizing. Set " + "'kv_cache_config.avg_seq_len' in the YAML configuration to the " + "workload's average total sequence length for an accurate KV/SSM " + "pool ratio.") + return fallback_capacity + def _get_quota_from_max_tokens(self, max_tokens: int) -> int: attention_quota = super()._get_quota_from_max_tokens(max_tokens) num_request_lineages = self._max_resident_sequences() + snapshot_slots = self._num_ssm_snapshots_for_capacity( + max_tokens, self.kv_cache_config) state_slots = (num_request_lineages + self._num_reserved_dummy_slots + - self._num_ssm_snapshots_for_capacity( - max_tokens, self.kv_cache_config)) + snapshot_slots) state_quota = state_slots * self._mamba_state_bytes_per_slot() # Once the plan contains any non-live SSM capacity, reserve one partial # attention page per request lineage. This remains conservative when @@ -2794,7 +2826,7 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int: extra_attention_quota = (num_request_lineages * self._attention_cache_bytes_per_token() * self.tokens_per_block - if state_slots > num_request_lineages else 0) + if snapshot_slots > 0 else 0) return attention_quota + state_quota + extra_attention_quota def _get_max_tokens_from_quota(self, quota: int) -> float: @@ -2817,20 +2849,6 @@ def _get_max_tokens_from_quota(self, quota: int) -> float: high = mid return low - def _planned_token_capacity( - self, - kv_cache_config: KvCacheConfig, - gpu_quota: int, - ) -> int: - capacity = self._get_max_tokens_from_quota(gpu_quota) - capacity = min( - capacity, - self.max_seq_len * self._max_resident_sequences(), - ) - if kv_cache_config.max_tokens is not None: - capacity = min(capacity, kv_cache_config.max_tokens) - return max(0, int(capacity)) - def _minimum_live_gpu_quota(self) -> int: """Return the minimum quota for live states and one attention page.""" attention_block_quota = (self._attention_cache_bytes_per_token() * @@ -2871,26 +2889,31 @@ def _build_cache_config( ], ) - num_sequences = self._max_resident_sequences() - planned_capacity = self._planned_token_capacity(kv_cache_config, - cache_tiers[0].quota) - capacity_per_request, capacity_remainder = divmod( - planned_capacity, num_sequences) - typical_ssm_slots = self._ssm_slots_per_request_for_typical_batch( - planned_capacity, kv_cache_config) + dummy_requests = [ + KVCacheDesc(capacity=0, history_length=0) + for _ in range(self._num_reserved_dummy_slots) + ] + constraints = [ + replace( + batch, + kv_caches=[*batch.kv_caches, *dummy_requests], + ) for batch in config.constraints + ] - typical_step = BatchDesc([ - KVCacheDesc( - capacity=capacity_per_request + int(i < capacity_remainder), - history_length=max( - 0, capacity_per_request + int(i < capacity_remainder) - 1), - num_ssm_slots=typical_ssm_slots[i], - ) for i in range(num_sequences) - ]) + typical_step = config.typical_step + if config.initial_pool_ratio is None: + typical_capacity = self._get_typical_request_capacity( + kv_cache_config) + request_descs = self._typical_request_descs(typical_capacity, + kv_cache_config) + typical_step = BatchDesc(request_descs * + self._max_resident_sequences() + + dummy_requests) return replace( config, layers=layers, typical_step=typical_step, + constraints=constraints, # SSM lifecycles require minimum-snapshot commit semantics. The # flag is harmless when reuse is disabled because no commits are # attempted, while the runtime config still needs the invariant. diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index dca766798c14..56e2e71d265c 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -422,6 +422,8 @@ def load_config_and_apply_defaults( config.pretrained_config) model_cls = AutoModelForCausalLM._resolve_class(config) + use_kv_cache_manager_v2 = ( + llm_args.kv_cache_config.use_kv_cache_manager_v2) # model_cls is None when the architecture is unknown/unsupported. model_defaults = {} @@ -435,16 +437,6 @@ def load_config_and_apply_defaults( f"Applied model defaults for {model_cls.__name__}: {applied_defaults}" ) - use_kv_cache_manager_v2 = llm_args.kv_cache_config.use_kv_cache_manager_v2 - _resolve_kv_cache_manager_v2_auto(llm_args, model_defaults) - _validate_and_adjust_mamba_snapshot_config(config, llm_args) - if use_kv_cache_manager_v2 == "auto": - logger.info( - "Resolved use_kv_cache_manager_v2='auto' to %s for %s", - llm_args.kv_cache_config.use_kv_cache_manager_v2, - model_cls.__name__ - if model_cls is not None else "unknown model") - # The transceiver preference follows the checkpoint's original # architecture: _resolve_class may rewrite it to an execution class # (e.g. MTPDraftModelForCausalLM), which must not drop the target @@ -458,6 +450,15 @@ def load_config_and_apply_defaults( # Resolve "auto" sentinel values after model defaults are applied. _resolve_transceiver_runtime_auto(llm_args, preference_cls, config.pretrained_config) + _resolve_kv_cache_manager_v2_auto( + llm_args, model_defaults, original_setting=use_kv_cache_manager_v2) + _validate_and_adjust_mamba_snapshot_config(config, llm_args) + if use_kv_cache_manager_v2 == "auto": + logger.info( + "Resolved use_kv_cache_manager_v2='auto' to %s for %s", + llm_args.kv_cache_config.use_kv_cache_manager_v2, + model_cls.__name__ + if model_cls is not None else "unknown model") return llm_args diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 7331a707ba39..394753181235 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3823,20 +3823,20 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): default=None, min_length=1, status="prototype", - description= - "Initial pool ratios for KV cache manager v2. When used by DeepSeek-V4, " - "values map to KVCacheManagerV2 pool_group_id order and must sum to 1.0. " - "When set, DeepSeek-V4 uses this directly and avg_seq_len does not take effect." - ) + description="Initial pool ratios for KV cache manager v2. Values map to " + "KVCacheManagerV2 pool_group_id order and must sum to 1.0. Hybrid Mamba " + "models and DeepSeek-V4 use this directly, so avg_seq_len does not take " + "effect when this is set.") # This is a pure python field, not a pybind field. It is only for the Pytorch backend. avg_seq_len: Optional[PositiveInt] = Field( default=None, status="prototype", description= - "Average sequence length used by DeepSeek-V4 to build the KV cache manager v2 " - "typical step. If unset, max_seq_len is used. This does not take effect when " - "pool_ratio is set.") + "Average total sequence length of the serving workload, used to build the " + "KV cache manager v2 typical step for hybrid Mamba models and DeepSeek-V4. " + "Hybrid Mamba models warn and fall back to half of max_seq_len when this is " + "unset. This does not take effect when pool_ratio is set.") # This is a pure python field, not a pybind field. It is only for the Pytorch backend. block_reuse_policy: Literal[ diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index 159ebbd7315a..63b360f5584b 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -558,9 +558,18 @@ def _compute_applied(defaults: Dict[str, Any], def _resolve_kv_cache_manager_v2_auto( - llm_args: 'TorchLlmArgs', model_defaults_dict: Dict[str, Any]) -> bool: - """Resolve the KV cache manager auto setting after model defaults are applied.""" - setting = llm_args.kv_cache_config.use_kv_cache_manager_v2 + llm_args: 'TorchLlmArgs', + model_defaults_dict: Dict[str, Any], + original_setting: Optional[Union[bool, str]] = None) -> bool: + """Resolve the KV cache manager auto setting after model defaults are applied. + + The transceiver runtime auto setting must be resolved first. In + disaggregated serving, hybrid Mamba V2 requires the Python transceiver with + NIXL, so an incompatible route falls back to V1 unless the user explicitly + selected V2. + """ + setting = (llm_args.kv_cache_config.use_kv_cache_manager_v2 + if original_setting is None else original_setting) if setting != "auto": return setting @@ -574,6 +583,18 @@ def _resolve_kv_cache_manager_v2_auto( "Model default kv_cache_config.use_kv_cache_manager_v2 must be " f"True, False, or 'auto', got {model_default!r}.") + transceiver_config = llm_args.cache_transceiver_config + if (model_default and transceiver_config is not None + and transceiver_config.backend is not None): + effective_backend, _ = transceiver_config._resolve_default_backend() + runtime = transceiver_config.transceiver_runtime + if effective_backend != "NIXL" or runtime != "PYTHON": + logger.info( + "KV cache manager V2 is the model default, but disaggregated " + "serving uses transceiver_runtime=%r with backend=%r; " + "falling back to V1.", runtime, effective_backend) + model_default = False + llm_args.kv_cache_config.use_kv_cache_manager_v2 = model_default return model_default diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index dee3b8022d8a..1a03c885b73a 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -168,7 +168,6 @@ LayerConfig = AttentionLayerConfig | SsmLayerConfig class KVCacheDesc: capacity: int history_length: int - num_ssm_slots: int = 1 @dataclass(slots=True) class BatchDesc: diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py index ed3912070985..8f374e7a37f7 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py @@ -145,11 +145,9 @@ def __post_init__(self) -> None: class KVCacheDesc: capacity: int history_length: int - num_ssm_slots: int = 1 def __post_init__(self) -> None: assert 0 <= self.history_length <= self.capacity - assert self.num_ssm_slots > 0 # A batch of requests, working as a use case the KVCacheManager must always support. diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py index d1aba2f540fd..303d79e9ee13 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py @@ -1077,6 +1077,10 @@ def plan_committed_block_drop(self) -> PlannedDropHandle | None: elif ssm_lc_id is not None: return None + def get_committed_page(block: Block, life_cycle_id: LifeCycleId) -> CommittedPage | None: + page_ref = block.storage[life_cycle_id] + return page_ref() if page_ref is not None else None + end = BlockOrdinal(len(matched_blocks)) pages_to_drop: list[CommittedPage] = [] for lc_idx, lc in self.manager._life_cycles.items(): @@ -1090,19 +1094,13 @@ def plan_committed_block_drop(self) -> PlannedDropHandle | None: window_start = min(stale_range.end, end) for ordinal in typed_range(window_start, end): tree_block = matched_blocks[ordinal] - page_ref = tree_block.storage[lc_idx] - if page_ref is None: - return None - page = page_ref() + page = get_committed_page(tree_block, lc_idx) if page is None: return None pages_to_drop.append(page) if ssm_lc_id is not None: - page_ref = matched_blocks[-1].storage[ssm_lc_id] - if page_ref is None: - return None - page = page_ref() + page = get_committed_page(matched_blocks[-1], ssm_lc_id) if not isinstance(page, SsmCommittedPage): return None pages_to_drop.append(page) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py index 22b0b70f55d3..69836ab0efcf 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py @@ -956,22 +956,12 @@ def _compute_slots_for_batch( """Compute the minimum number of slots per pool group to support a BatchDesc.""" num_slots = filled_list(0, self.num_pool_groups) ssm_lc_idx = self._life_cycles.ssm_life_cycle_id - total_ssm_slots = sum(kv.num_ssm_slots for kv in batch.kv_caches) - num_request_lineages = len(batch.kv_caches) - # An extra SSM slot indicates that this batch can retain snapshots. - # Snapshot alignment is unknown while storage is being sized, so once - # such capacity exists, reserve one partial attention page per request - # lineage. V2 replaces covered partial pages within a lineage, making N - # a safe upper bound. This does not require S >= 2N: when N < S < 2N, - # the bound merely over-reserves attention pages. The guard keeps the - # default num_ssm_slots=1 descriptors used by non-Mamba managers from - # paying this Mamba-specific overhead. - has_extra_ssm_slot = total_ssm_slots > num_request_lineages sys_blocks = batch.system_prompt_length // tokens_per_block for lc_idx, lc in typed_enumerate(self._life_cycles.get()): pg_idx = self.get_pool_group_index(lc_idx) if lc_idx == ssm_lc_idx: - num_slots[pg_idx] += total_ssm_slots + # SSM: always 1 dedicated block per request, never shared. + num_slots[pg_idx] += len(batch.kv_caches) continue # Shared sys blocks (counted once): union of non-stale sys blocks across all requests. # A sys block needs memory if it's non-stale for ANY request. @@ -1008,11 +998,6 @@ def _compute_slots_for_batch( ) else: num_slots[pg_idx] += unique_non_stale - # Retained partial-page copies are additional physical blocks, not - # tokens in KVCacheDesc.capacity, so add their safe upper bound - # after token staleness and scratch-sharing calculations. - if has_extra_ssm_slot: - num_slots[pg_idx] += num_request_lineages return num_slots def _slots_to_bytes( diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index c1fbdb9b1bc2..7bb1457f2bed 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -6291,6 +6291,9 @@ def test_fp8(self, enable_block_reuse, mocker): kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.8, enable_block_reuse=enable_block_reuse) + if enable_block_reuse: + kv_cache_config.avg_seq_len = 2048 + kv_cache_config.mamba_state_config.periodic_snapshot_interval = 256 # DeepGEMM MoE kernels only support datacenter Blackwell (SM100/SM103). # Fall back to the CUTLASS MoE backend (which supports FP8 block scales) # on other architectures such as Hopper (SM90) and consumer Blackwell @@ -6444,6 +6447,9 @@ def test_nvfp4(self, tp_size, ep_size, attention_dp, moe_backend, kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.9, enable_block_reuse=enable_block_reuse) + if enable_block_reuse: + kv_cache_config.avg_seq_len = 2048 + kv_cache_config.mamba_state_config.periodic_snapshot_interval = 256 cuda_graph_config = CudaGraphConfig(max_batch_size=256, enable_padding=True) @@ -7154,6 +7160,7 @@ def test_nvfp4_4gpus_block_reuse(self, tp_size, ep_size, f"{llm_models_root()}/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4", kv_cache_config=KvCacheConfig( enable_block_reuse=True, + avg_seq_len=2048, mamba_ssm_cache_dtype="float16", mamba_state_config=MambaStateConfig( periodic_snapshot_interval=periodic_snapshot_interval), @@ -7229,7 +7236,10 @@ def test_nvfp4_8gpus_mtp(self): model_path, kv_cache_config=KvCacheConfig( enable_block_reuse=True, + avg_seq_len=2048, mamba_ssm_cache_dtype="float16", + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=256), free_gpu_memory_fraction=0.5, ), max_batch_size=32, @@ -7583,6 +7593,7 @@ def test_nvfp4_4gpus_block_reuse(self, tp_size, ep_size, f"{llm_models_root()}/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4", kv_cache_config=KvCacheConfig( enable_block_reuse=True, + avg_seq_len=2048, mamba_ssm_cache_dtype="float16", mamba_state_config=MambaStateConfig( periodic_snapshot_interval=periodic_snapshot_interval), diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index b2728622667e..e7d61248e7b4 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -534,8 +534,8 @@ def test_hybrid_models_default_to_v2_and_python_transceiver(monkeypatch): ) model_defaults = model_cls.get_model_defaults(llm_args) apply_model_defaults_to_llm_args(llm_args, model_defaults) - _resolve_kv_cache_manager_v2_auto(llm_args, model_defaults) _resolve_transceiver_runtime_auto(llm_args, model_cls) + _resolve_kv_cache_manager_v2_auto(llm_args, model_defaults, original_setting="auto") assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True assert llm_args.kv_cache_config.enable_block_reuse is False assert llm_args.cache_transceiver_config.transceiver_runtime == "PYTHON" @@ -1001,14 +1001,20 @@ def test_v2_hybrid_snapshot_sizing_scales_with_pp_and_explicit_rules(): ) assert mgr._num_ssm_snapshots_for_capacity(512, config) == 24 - assert mgr._ssm_slots_per_request_for_typical_batch(512, config) == [4] * 8 + assert mgr._num_ssm_states_per_typical_request(128, config) == 4 + assert [desc.capacity for desc in mgr._typical_request_descs(128, config)] == [ + 32, + 32, + 32, + 32, + ] periodic_config = KvCacheConfig( enable_block_reuse=True, mamba_state_config=MambaStateConfig(periodic_snapshot_interval=48), ) - assert mgr._ssm_slots_per_request_for_typical_batch(47, periodic_config) == [1] * 8 - assert mgr._ssm_slots_per_request_for_typical_batch(48, periodic_config) == [2] + [1] * 7 + assert mgr._num_ssm_states_per_typical_request(47, periodic_config) == 1 + assert mgr._num_ssm_states_per_typical_request(48, periodic_config) == 1 def test_cpp_mamba_estimator_handles_disabled_snapshots_without_attention(): @@ -1087,10 +1093,7 @@ def test_hybrid_mtp_layout_honors_explicit_base_partition(): kv_cache_config=KvCacheConfig(enable_block_reuse=False), spec_config=spec_config, ) - extra_attention_bound = ( - 4096 * (rank + 1) if manager_cls is MambaHybridCacheManagerV2 else 0 - ) - assert cache_cost == (64 * (rank + 1), 2400 + extra_attention_bound) + assert cache_cost == (64 * (rank + 1), 2400) def test_hybrid_separate_mtp_draft_estimator_has_no_mamba_state(): @@ -1115,7 +1118,7 @@ def test_hybrid_separate_mtp_draft_estimator_has_no_mamba_state(): spec_config=spec_config, use_separate_draft_kv_cache=True, ) - assert target_cost == (64, 6496) + assert target_cost == (64, 2400) _, local_mamba_layers, local_attention_layers = _get_local_mamba_cache_layout( model_config, @@ -1135,25 +1138,25 @@ def test_hybrid_separate_mtp_draft_estimator_has_no_mamba_state(): spec_config=spec_config, is_draft=True, ) - assert draft_cost == (64 * rank, 4096 * rank) + assert draft_cost == (64 * rank, 0) @pytest.mark.parametrize( ("spec_config", "enable_attention_dp", "expected_intercept"), [ - (None, False, 1728), - (MTPDecodingConfig(max_draft_len=4), True, 1792), + (None, False, 320), + (MTPDecodingConfig(max_draft_len=4), True, 384), ( MTPDecodingConfig( max_draft_len=4, draft_len_schedule={1: 4, 2: 2, 3: 1}, ), True, - 1984, + 576, ), ], ) -def test_v2_hybrid_estimator_accounts_for_ssm_slot_attention_bound( +def test_v2_hybrid_estimator_counts_dummy_states_without_attention_capacity( monkeypatch, spec_config, enable_attention_dp, expected_intercept ): monkeypatch.setattr( @@ -1208,18 +1211,6 @@ def estimate(interval): assert aligned[1] == unaligned[1] -def test_v2_hybrid_planned_capacity_is_bounded_by_resident_sequences(): - mgr = object.__new__(MambaHybridCacheManagerV2) - mgr.max_batch_size = 4 - mgr.mapping = Mapping(world_size=2, rank=0, tp_size=1, pp_size=2) - mgr.max_seq_len = 128 - mgr._get_max_tokens_from_quota = lambda quota: 1 << 30 - - capacity = mgr._planned_token_capacity(KvCacheConfig(max_tokens=None), 1 << 40) - - assert capacity == 4 * 2 * 128 - - def _base_attention_layer_configs(num_layers): return [ AttentionLayerConfig( @@ -1230,7 +1221,7 @@ def _base_attention_layer_configs(num_layers): ] -def test_v2_hybrid_typical_batch_uses_num_ssm_slots(): +def test_v2_hybrid_typical_batch_splits_capacity_across_ssm_states_and_dummies(): mgr = object.__new__(MambaHybridCacheManagerV2) mgr.kv_cache_type = CacheTypeCpp.SELF mgr.head_dim_per_layer = [64, 64] @@ -1253,10 +1244,14 @@ def test_v2_hybrid_typical_batch_uses_num_ssm_slots(): mgr.num_extra_kv_tokens = 0 mgr.get_layer_bytes_per_token = lambda **kwargs: 8 mgr._minimum_live_gpu_quota = lambda: 0 - mgr._planned_token_capacity = lambda *_args: 128 kv_cache_config = KvCacheConfig( enable_partial_reuse=True, - mamba_state_config=MambaStateConfig(periodic_snapshot_interval=48), + avg_seq_len=96, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=48, + additional_snapshot_offsets_from_start=[32], + additional_snapshot_offsets_from_end=[0], + ), ) mgr.kv_cache_config = kv_cache_config constraints = [BatchDesc([KVCacheDesc(capacity=64, history_length=0)])] @@ -1272,12 +1267,38 @@ def test_v2_hybrid_typical_batch_uses_num_ssm_slots(): assert isinstance(config.layers[0], SsmLayerConfig) assert config.layers[1] is base_layers[1] - assert len(config.typical_step.kv_caches) == 2 - assert [kv.num_ssm_slots for kv in config.typical_step.kv_caches] == [3, 2] - assert config.constraints is constraints - assert not hasattr(config.typical_step, "ssm_cache") - assert not hasattr(config.typical_step, "additional_attention_cache") - assert not hasattr(config.typical_step.kv_caches[0], "num_extra_attention_slots") + assert config.typical_step == BatchDesc( + [KVCacheDesc(capacity=32, history_length=31)] * 6 + + [KVCacheDesc(capacity=0, history_length=0)] + ) + assert config.constraints == [ + BatchDesc( + [ + KVCacheDesc(capacity=64, history_length=0), + KVCacheDesc(capacity=0, history_length=0), + ] + ) + ] + assert sum(kv.capacity for kv in config.typical_step.kv_caches) == 2 * 96 + assert not hasattr(config.typical_step.kv_caches[0], "num_ssm_slots") + + +def test_v2_hybrid_warns_when_avg_seq_len_is_missing(monkeypatch): + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.max_seq_len = 4096 + warnings_seen = [] + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.mamba_cache_manager.logger.warning", + lambda message: warnings_seen.append(message), + ) + + capacity = mgr._get_typical_request_capacity(KvCacheConfig()) + + assert capacity == 2048 + assert len(warnings_seen) == 1 + assert "kv_cache_config.avg_seq_len" in warnings_seen[0] + assert "max_seq_len / 2=2048" in warnings_seen[0] + assert "workload's average total sequence length" in warnings_seen[0] def test_v2_hybrid_rejects_quota_below_live_state_floor(): @@ -2218,58 +2239,6 @@ def test_hybrid_replay_buffers_size_by_tokens_per_gen_step(builder): mgr.shutdown() -@skip_no_cuda -def test_v2_hybrid_replay_bookkeeping_matches_checkpoint_predicate(monkeypatch): - spec_config = _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5) - mgr = _build_v2_hybrid_with_mamba_layer( - max_batch_size=4, - spec_config=spec_config, - use_replay_state_update=True, - ) - monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.mamba_cache_manager._promote_mamba_state_triton", - lambda *args, **kwargs: None, - ) - try: - # Derive the checkpoint threshold from the manager's own replay - # metadata instead of hardcoding it: replay_history_size is - # max(MIN_REPLAY_HISTORY_SIZE, tokens_per_gen_step), so the boundary - # is not simply tokens_per_gen_step. - replay_metadata = mgr.get_replay_state_update_metadata() - step_width = replay_metadata.replay_step_width - history_size = replay_metadata.replay_history_size - - slot = torch.tensor([0], dtype=torch.int32, device="cuda") - attn_metadata = SimpleNamespace(num_seqs=1, num_contexts=0) - - def advance(accepted): - mgr.update_mamba_states( - attn_metadata, - torch.tensor([accepted], dtype=torch.int32, device="cuda"), - state_indices=slot, - ) - - # Accumulate below the checkpoint threshold: while - # prev + step_width <= history_size the manager keeps writing into the - # same cache_buf_idx and grows prev_num_accepted_tokens monotonically. - prev = 0 - while prev + step_width <= history_size: - advance(1) - prev += 1 - assert mgr.prev_num_accepted_tokens[0].item() == prev - assert mgr.cache_buf_idx[0].item() == 0 - - # The next step crosses the threshold (prev + step_width > history_size): - # the manager starts a fresh checkpoint, resetting - # prev_num_accepted_tokens to the accepted count and flipping - # cache_buf_idx to the other buffer. - advance(2) - assert mgr.prev_num_accepted_tokens[0].item() == 2 - assert mgr.cache_buf_idx[0].item() == 1 - finally: - mgr.shutdown() - - def test_v2_hybrid_replay_update_skips_dummy_and_padding_rows(monkeypatch): mgr = object.__new__(MambaHybridCacheManagerV2) mgr.local_num_mamba_layers = 1 diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index 36a49e73f8ec..469c66a61f3f 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -2764,94 +2764,23 @@ def test_typical_step_long_sequences(self): self.assertLess(ratio[0], 0.15) manager.shutdown() - def test_num_ssm_slots_controls_ssm_slot_count(self): - """SSM sizing uses explicit state-slot count, not token capacity.""" - manager = KVCacheManager(self._make_hybrid_config()) - ssm_lc = manager._life_cycles.ssm_life_cycle_id - assert ssm_lc is not None - ssm_pg = manager._storage.get_pool_group_index(ssm_lc) - - batch = BatchDesc( - kv_caches=[ - KVCacheDesc( - capacity=1024, - history_length=1023, - num_ssm_slots=3, - ) - ] - * 2 - ) - slots = manager._storage._compute_slots_for_batch(batch, self.TOKENS_PER_BLOCK, None) - self.assertEqual(slots[ssm_pg], 6) - - default_batch = BatchDesc(kv_caches=[KVCacheDesc(capacity=4096, history_length=4095)] * 2) - default_slots = manager._storage._compute_slots_for_batch( - default_batch, self.TOKENS_PER_BLOCK, None - ) - self.assertEqual(default_slots[ssm_pg], 2) - manager.shutdown() - - def test_num_ssm_slots_must_be_positive(self): - with self.assertRaises(AssertionError): - KVCacheDesc(capacity=1, history_length=0, num_ssm_slots=0) - - def test_ssm_slots_bound_additional_attention_capacity(self): + def test_zero_capacity_request_reserves_only_an_ssm_slot(self): + """Every request reserves one SSM slot, including a zero-token dummy.""" manager = KVCacheManager(self._make_hybrid_config()) ssm_lc = manager._life_cycles.ssm_life_cycle_id assert ssm_lc is not None ssm_pg = manager._storage.get_pool_group_index(ssm_lc) attn_pg = 1 - ssm_pg - batch = BatchDesc([KVCacheDesc(capacity=64, history_length=63, num_ssm_slots=3)]) - slots = manager._storage._compute_slots_for_batch(batch, self.TOKENS_PER_BLOCK, None) - - # Two SSM slots beyond the live state imply at most one retained - # partial attention page for this request lineage. - self.assertEqual(slots[attn_pg], 3) - manager.shutdown() - - def test_ssm_slot_attention_bound_applies_to_each_lifecycle(self): - config = self._make_config(enable_swa_scratch_reuse=True) - manager = KVCacheManager(config) - base_batch = BatchDesc([KVCacheDesc(capacity=512, history_length=32)]) - snapshot_batch = BatchDesc([KVCacheDesc(capacity=512, history_length=32, num_ssm_slots=3)]) - - base_slots = manager._storage._compute_slots_for_batch( - base_batch, self.TOKENS_PER_BLOCK, config.swa_scratch_reuse - ) - snapshot_slots = manager._storage._compute_slots_for_batch( - snapshot_batch, self.TOKENS_PER_BLOCK, config.swa_scratch_reuse - ) - expected_increase = [0] * manager._storage.num_pool_groups - for life_cycle_id, _ in manager._life_cycles.attention_life_cycles(): - pool_group = manager._storage.get_pool_group_index(life_cycle_id) - expected_increase[pool_group] += 1 - - self.assertEqual( - [actual - base for actual, base in zip(snapshot_slots, base_slots)], - expected_increase, - ) - manager.shutdown() - - def test_ssm_slot_attention_bound_reserves_each_lineage(self): - manager = KVCacheManager(self._make_hybrid_config()) - ssm_lc = manager._life_cycles.ssm_life_cycle_id - assert ssm_lc is not None - ssm_pg = manager._storage.get_pool_group_index(ssm_lc) - attn_pg = 1 - ssm_pg batch = BatchDesc( - [ - KVCacheDesc(capacity=64, history_length=63, num_ssm_slots=3), - *[KVCacheDesc(capacity=64, history_length=63)] * 3, + kv_caches=[ + KVCacheDesc(capacity=64, history_length=63), + KVCacheDesc(capacity=0, history_length=0), ] ) - slots = manager._storage._compute_slots_for_batch(batch, self.TOKENS_PER_BLOCK, None) - - self.assertEqual(slots[ssm_pg], 6) - # Any extra SSM slot enables the conservative upper bound of one - # partial attention page for every request lineage. - self.assertEqual(slots[attn_pg], 12) + self.assertEqual(slots[ssm_pg], 2) + self.assertEqual(slots[attn_pg], 2) manager.shutdown() def test_constraints_floor_typical_step(self): diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 1ce215981460..183481805aed 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -494,7 +494,9 @@ def test_kv_cache_manager_v2_auto_uses_model_default(self, explicit_auto): model_defaults = {"kv_cache_config": {"use_kv_cache_manager_v2": True}} apply_model_defaults_to_llm_args(llm_args, model_defaults) - _resolve_kv_cache_manager_v2_auto(llm_args, model_defaults) + _resolve_kv_cache_manager_v2_auto(llm_args, + model_defaults, + original_setting="auto") assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True @@ -505,6 +507,45 @@ def test_kv_cache_manager_v2_auto_falls_back_to_false(self): assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False + @pytest.mark.parametrize( + ("backend", "runtime"), + [ + ("NIXL", "CPP"), + ("UCX", None), + ("MPI", None), + ], + ) + def test_kv_cache_manager_v2_auto_falls_back_for_incompatible_disagg( + self, backend, runtime): + llm_args = TorchLlmArgs( + model="/tmp/dummy_model", + cache_transceiver_config=CacheTransceiverConfig( + backend=backend, transceiver_runtime=runtime), + ) + model_defaults = {"kv_cache_config": {"use_kv_cache_manager_v2": True}} + + apply_model_defaults_to_llm_args(llm_args, model_defaults) + _resolve_kv_cache_manager_v2_auto(llm_args, + model_defaults, + original_setting="auto") + + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False + + def test_kv_cache_manager_v2_auto_keeps_python_nixl_model_default(self): + llm_args = TorchLlmArgs( + model="/tmp/dummy_model", + cache_transceiver_config=CacheTransceiverConfig( + backend="NIXL", transceiver_runtime="PYTHON"), + ) + model_defaults = {"kv_cache_config": {"use_kv_cache_manager_v2": True}} + + apply_model_defaults_to_llm_args(llm_args, model_defaults) + _resolve_kv_cache_manager_v2_auto(llm_args, + model_defaults, + original_setting="auto") + + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True + @pytest.mark.parametrize("user_setting", [False, True]) def test_kv_cache_manager_v2_explicit_value_overrides_model_default( self, user_setting): @@ -761,15 +802,6 @@ def test_KvCacheConfig_requires_v2_for_additional_snapshot_offsets( assert getattr(config.mamba_state_config, field) == offsets -def test_KvCacheConfig_allows_periodic_snapshots_with_v1(): - config = KvCacheConfig( - mamba_state_config=MambaStateConfig(periodic_snapshot_interval=64), - use_kv_cache_manager_v2=False, - ) - - assert config.mamba_state_config.periodic_snapshot_interval == 64 - - def test_KvCacheConfig_migrates_deprecated_mamba_interval(monkeypatch): warnings_seen = [] monkeypatch.setattr(llm_args_mod.logger, "warning", @@ -3501,13 +3533,32 @@ def test_default_backend_resolves_to_nixl_and_adopts_preference( # creation time. assert args.cache_transceiver_config.backend == "DEFAULT" - def test_default_backend_env_override_falls_back_to_cpp(self, monkeypatch): - """DEFAULT + TRTLLM_USE_UCX_KVCACHE=1 means effective UCX -> C++.""" - monkeypatch.delenv("TRTLLM_USE_NIXL_KVCACHE", raising=False) - monkeypatch.setenv("TRTLLM_USE_UCX_KVCACHE", "1") + @pytest.mark.parametrize( + "backend_env", + ["TRTLLM_USE_UCX_KVCACHE", "TRTLLM_USE_MPI_KVCACHE"], + ) + def test_default_backend_env_override_falls_back_to_v1_cpp( + self, monkeypatch, backend_env): + """An incompatible DEFAULT route falls back to the V1 C++ path.""" + for env_var in ( + "TRTLLM_USE_NIXL_KVCACHE", + "TRTLLM_USE_UCX_KVCACHE", + "TRTLLM_USE_MOONCAKE_KVCACHE", + "TRTLLM_USE_MPI_KVCACHE", + ): + monkeypatch.delenv(env_var, raising=False) + monkeypatch.setenv(backend_env, "1") args = self._disagg_args(backend="DEFAULT") + model_defaults = {"kv_cache_config": {"use_kv_cache_manager_v2": True}} + apply_model_defaults_to_llm_args(args, model_defaults) + _resolve_transceiver_runtime_auto(args, _PreferPythonTransceiverModel) + _resolve_kv_cache_manager_v2_auto(args, + model_defaults, + original_setting="auto") + assert args.cache_transceiver_config.transceiver_runtime is None + assert args.kv_cache_config.use_kv_cache_manager_v2 is False def test_disagg_disabled_is_noop(self): """Resolver never creates a config when cache_transceiver_config is None."""