[None][feat] Add V2 Mamba hybrid cache with snapshot reuse

[VALLIS-NERIA]

@coderabbitai summary

Stack

• Base branch: agent/mamba-scheduling-disagg-lifecycle
• Prerequisite: NVIDIA/TensorRT-LLM#16427
• This PR contains exactly one signed-off commit relative to its base.

Description

This PR adds V2 Mamba hybrid-cache snapshot reuse on top of the scheduler and
resource-manager contracts in the prerequisite PR.

The changes:

• Keep the V1 C++ Mamba manager as the default and select V2 only when
  use_kv_cache_manager_v2: true is explicitly configured. V2 disaggregated
  serving and V1 additional fixed snapshot offsets are rejected up front.
• Move snapshot policy under mamba_state_config, with
  periodic_snapshot_interval, additional_snapshot_offsets_from_start, and
  additional_snapshot_offsets_from_end. The old
  mamba_state_cache_interval input/property remains as a deprecated
  compatibility alias.
• Add V2MambaHybridCacheManager, storing attention KV blocks and Mamba
  convolution/SSM state in role-aware KVCMv2 pools while keeping generic
  KVCMv2 policy-free.
• Use expected_snapshot_points as the snapshot-boundary source and support
  periodic, fixed-from-start, and fixed-from-end boundaries, including the
  final prompt boundary.
• Keep CUDA state-index addresses stable across dummy insertion, mirror replay
  checkpoint bookkeeping, clean stale request-index mappings, and support
  attention-only and Mamba-only PP ranks.
• Match model PP semantics under MTP: partition base layers first, append spec
  attention layers to the last PP rank, and use the same exact masks for cache
  construction and affine memory sizing. Separate attention-only draft caches
  no longer receive phantom Mamba fixed cost.
• Remove the PR-specific prefix-sibling mechanism. Upstream
  commit_min_snapshot plus the corrected radix root-replacement order gives
  safe reuse; only the latest same-block partial snapshot is retained.
• Size partial-snapshot attention storage to at most one copied page per live
  request and scan every physical attention/state buffer in invalid-value
  diagnostics.

The detailed design decisions, experiment results, and verification record are
in PR3_MAMBA_V2_REVIEW.md.

Test Coverage

• Full Mamba cache-manager suite: 85 passed.
• Combined Mamba, cache-budget split, dual-pool, and executor-creator suites:
  151 passed.
• Full LLM-args suite: 255 passed, 3 skipped; a preceding local attempt hit
  an unrelated MPI worker-start timeout before any Mamba test logic.
• Generic KVCMv2 and executor-creator focused suites: 24 passed.
• Runtime KVCMv2 SSM support: 13 passed.
• Runtime KVCMv2 no-batching coverage: 24 passed, 12 skipped.
• Golden manifest check, Python syntax checks, diff checks, and all changed-file
  pre-commit hooks passed.

PR Checklist

• Please check this after reviewing the above items as appropriate for this PR.

[VALLIS-NERIA]

Focused cleanup review of the PR3-only diff against the latest PR1 base.

[VALLIS-NERIA]

[P1] This forwards enable_swa_scratch_reuse into a cache containing SsmLayerConfigs, but KVCacheManagerV2._prepare_page_table_tensor() still hard-codes Role.KEY / Role.KEY_BLOCK_SCALE for every layer in its scratch-reuse branch. Mamba layers register only SSM_STATE and CONV_STATE, so enabling SWA scratch reuse reaches get_mem_pool_base_address(mamba_layer, Role.KEY, PER_LAYER) and fails during construction. Please either reject/disable this combination explicitly for V2 Mamba, or generalize the per-layer scratch page-table branch with the same role hooks used by the non-scratch branch, and add a construction test for this configuration.

[VALLIS-NERIA]

Agreed. Generalized SWA scratch pointer and converter construction through the pool role hooks, preserved the SSM placeholder row, and made state views request the explicit per-layer base with PageIndexMode.SHARED. Added a constructor plus copy_batch_block_offsets GPU regression test. Addressed in 9b6d32c.

[VALLIS-NERIA]

[P2] This field is consumed as the total SSM slot count, not the number of reusable snapshots: the V2 builder adds max_batch_size live-state slots to snapshot slots before populating it, and None maps to one live slot in _compute_slots_for_batch(). The name therefore invites callers to omit the live slot, and the current >= 0 validation even permits an SSM request to plan zero slots. Please rename it to e.g. num_ssm_slots (default 1, require > 0), or keep snapshot semantics and have the storage planner add one live slot per request. Update the stub and tests with the same invariant.

[VALLIS-NERIA]

Agreed. Renamed the field to num_ssm_slots, made its default explicitly 1, and reject non-positive values. Runtime config, stub, storage sizing, and tests now use the same total-slot semantics. Addressed in 9b6d32c.

[VALLIS-NERIA]

[cleanup] The routing matrix still contains several leftovers. The earlier if is_disagg: returns on every path, so and not is_disagg at line 136 is redundant; default_cls is already V2MambaHybridCacheManager, so this block-reuse branch returns exactly the same class as the final return; and the V2 preference warning says it overrides a default that is already V2. Please collapse these branches. While doing so, remove the two test-only TRTLLM_USE_CPP_MAMBA delenv calls (production has no reader left) and update use_py_mamba_cache_manager() / MambaHybridCacheManager / Cpp-manager docstrings that still describe Cpp as the default or list only Mixed/Cpp.

[VALLIS-NERIA]

Agreed. Removed the redundant routing branches and stale TRTLLM_USE_CPP_MAMBA test cleanup, stopped warning when V2 is explicitly selected because it is already the default, and refreshed the manager-selection documentation. Addressed in 9b6d32c.

[VALLIS-NERIA]

[cleanup] _build_pool_mapping_tensors() is dead code: a repository-wide call-site search finds only this definition (plus an unrelated sparse-manager definition), while KVCacheManagerV2.__init__() calls _prepare_page_table_tensor() directly. The base implementation was already generalized in this PR through _get_pool_page_index_role() / _get_pool_paired_role(), so this 70-line duplicate is stale and bypasses the base implementation's exact_div checks. Please remove the whole method.

[VALLIS-NERIA]

Agreed. _build_pool_mapping_tensors had no callers after the base V2 page-table path became authoritative, so the entire dead implementation was removed. Addressed in 9b6d32c.

[VALLIS-NERIA]

[cleanup] These two overrides add no behavior: flush_state_transfers() is an unconditional no-op and its only caller already uses getattr, while update_resources() forwards the identical arguments to the inherited KVCacheManagerV2 method. Please remove both. The same initialization pass can also drop _full_attention_layer_mask / _combined_layer_mask (write-only) and the duplicate self.kv_cache_config assignment already performed by the base class.

[VALLIS-NERIA]

Agreed. Removed the no-op flush_state_transfers override, the pass-through update_resources override, write-only mask/config state, and the two one-use request helpers. The commit/snapshot decisions now live at their KVCacheManagerV2 call sites. Addressed in 9b6d32c.

[VALLIS-NERIA]

[cleanup/scope] model_type is used here only to retain conv_section_dims. The sole production consumer of that field is the disaggregated kv_extractor, which currently accesses the legacy facade's manager._impl; this PR's routing explicitly never selects V2 for disaggregated serving. Since V2 disagg is a later PR, please move this future-only field/model-specific branching there and keep the aggregate V2 core independent of transfer-layout metadata.

[VALLIS-NERIA]

Agreed. Removed model_type and conv_section_dims from V2 because they only serve the later disaggregated transport work. The factory now supplies model_type only to the legacy Mixed/Cpp implementations that consume it. Addressed in 9b6d32c.