[VALLIS-NERIA]

Description

Problem

Silent state corruption: with KV cache manager V2 (MambaHybridCacheManagerV2) and the Mamba2 or GDN replay state update (speculative decoding, use_replay_state_update=True), suspending a generation request corrupts its recurrent state.

The per-slot replay state lives outside the V2 pools, in plain torch tensors indexed by the SSM slot:

• prev_num_accepted_tokens (PNAT)
• cache_buf_idx
• old_x / old_B / old_dt / old_dA_cumsum, shaped [layers, num_ssm_slots, 2, ...]

A decoding request's true recurrent state is the V2 SSM checkpoint plus the replay history history[cache_buf_idx][0:PNAT].

A generation request is suspended in two places:

• the scheduler: V2 _try_evict_for_gen -> kv_cache_manager.suspend_request
• PyExecutor._rebalance_kv_pools_now: suspend all -> impl.adjust() -> resume all

While the request is suspended, V2 may move its SSM page to the host tier and give the slot to another cache. On resume the page can come back in a different slot:

• The SSM/conv checkpoint follows the page.
• The replay history stays at the old slot index, where the slot's new owner may already have reset or overwritten it.

So the request resumes with whatever replay state its new slot holds: PNAT=0 in the repro, or a previous owner's stale history. KDA (Kimi) already handles this: it commits the pending drafts at suspend and re-seeds at resume. Mamba2/GDN had no equivalent.

Fix

The changes are in MambaHybridCacheManagerV2 and apply to Mamba2/GDN replay only. The KDA and non-replay paths are unchanged.

• suspend_request: for a non-dummy request that holds a slot, copy the slot's replay state to new pinned host tensors (async D2H). The copy covers PNAT, cache_buf_idx and old_x / old_B / old_dt / old_dA_cumsum, including both halves of the double-buffered history. The copies are kept per request id. Then forget the request's state-index entry, as KDA does, and call the base suspend.

• _setup_state_indices:

  • If a suspended request comes back as a generation request, copy the saved state (async H2D) into whatever slot it now owns.
  • If it comes back as a context request (re-prefill), drop the saved state.

• free_resources: drops the saved state too.

• Stream ordering mirrors the KDA path (_on_kv_cache_stream):

  1. The KV-cache stream waits on the current stream.
  2. The copies run on the KV-cache stream.
  3. The current stream waits on the KV-cache stream.

  This orders the D2H after the request's last verify. It also orders it before anything a later owner of the slot writes: a context reset, a verify, or another request's restore. The H2D runs before the next forward.

• Exactness: it is a pure copy, so the resumed state is bit-exact.

• Dummy / padding requests are skipped. Their replay state is never read, and they keep their slot mapping.

• The Philox seed buffer (mamba_ssm_rand_seed) is intentionally not carried. The Mamba2 mixer passes only element 0 (rand_seed[:1]) to the kernels and advances the whole buffer on every forward. So the buffer acts as one shared Philox counter, not per-request state. Restoring an old per-slot value into slot 0 would rewind that shared counter.

• Size: each copy is the whole slot slice, which keeps the copy sync-free (no host read of PNAT or cache_buf_idx). For a large Nemotron-H this is on the order of tens of MB of pinned host memory per suspended request. Copying only the live cache_buf_idx half would halve it; that is left out to keep this change small.

Rebalance decision: config-time guard

_build_cache_config now rejects kv_cache_config.enable_kv_pool_rebalance=True with Mamba2/GDN replay on V2. This mirrors the existing KDA guard. The reasons:

• The replay tensors are allocated once, for the SSM pool's slot count at construction (_allocate_pool_replay_buffers).
• The replay kernels index them by slot without a bound.
• KvCacheManager::getPageIndexUpperBound is derived from the pool group's current hot-level slot count.
• adjust() -> _adjustLevel -> adjustCacheLevel re-targets the pool-group ratios.

So after a rebalance the SSM pool can hold more slots than the replay tensors cover, and a resumed request can land past their end. Rejecting the combination is safer than resizing these tensors in a short-term fix.

enable_kv_pool_rebalance is an opt-in prototype (default False). The error message points to enable_kv_pool_rebalance=False, or to TRTLLM_USE_MAMBA_REPLAY=0 for the non-replay path. Speculative decoding without replay can still use the rebalance.

Scope

This is a short-term safety net. The long-term direction is to hold the replay buffers in a KVCM PRIVATE SSM lifecycle (see #21088). They would then move with the slot like the SSM/conv states do, and this evacuation would no longer be needed.

Test Coverage

New tests in tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py:

• test_v2_replay_suspend_resume_keeps_the_replay_state[{slot_taken,slot_kept}-{mamba2_replay,gdn_replay}]
  • Uses a real V2 manager with large SSM slots, mirroring the KDA helpers _v2_kda_replay_manager_with_large_slots and _take_ssm_slot.
  • The request gets a distinctive checkpoint and replay state: PNAT=5, cache_buf_idx=1, random histories.
  • The request is suspended. In slot_taken, its slot goes to another cache whose writes overwrite it.
  • After resume and _setup_state_indices, the test asserts that:
    • the slot changed (taken case);
    • the SSM/conv checkpoint and every replay field equal the pre-suspend values, bit for bit, in the new slot;
    • the new owner's slot was left alone.
• test_v2_replay_suspended_request_resumed_as_context_starts_over[...]: a suspended request that comes back as a context request gets no stale history restored, either then or at its next generation step.
• test_v2_replay_suspend_keeps_nothing_for_freed_or_dummy_requests[...]: freeing a suspended request drops its saved state. A dummy's suspend saves nothing and keeps its slot mapping.
• test_v2_replay_refuses_kv_pool_rebalance[{replay,no_replay}-{mamba2,gdn}]: covers the new guard.

How the tests were run

• Hardware: one B200 on umbriel-b200-043.
• Image: tensorrt-llm:pytorch-26.09-py3-x86_64-ubuntu24.04-skip-tritondevel-202610230003-19760.
• Build: the main CI wheel tensorrt_llm-1.4.0rc2+ngcpytorch2609, with this branch's git-tracked Python sources overlaid on top. This is a Python-only change on base 011316470c.

Run	Result
New tests, base mamba_cache_manager.py (0113164)	8 failed, 4 passed
New tests, this PR	12 passed
Whole test_mamba_cache_manager.py, this PR	277 passed
test_kv_cache_v2_scheduler.py + test_kv_pool_rebalance.py, this PR	410 passed

On the base, these 8 tests fail:

• test_v2_replay_suspend_resume_keeps_the_replay_state[slot_taken-mamba2_replay] and [slot_taken-gdn_replay]: AssertionError: prev_num_accepted_tokens. The resumed slot has lost the history; this is the bug.
• test_v2_replay_refuses_kv_pool_rebalance[replay-*]: no guard, so nothing is raised.
• The context and freed/dummy tests, for both layouts. They check the new saved-state bookkeeping, which the base does not have.

The 4 passing on the base are the slot_kept and no_replay cases.

The existing KDA suspend tests, including test_v2_kda_replay_suspend_resume_matches_an_unsuspended_run, still pass with this change.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

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

🤖 Generated with Claude Code