[VALLIS-NERIA]

Description

Stacked on #21088 (layer and buffer reuse policies). Until that merges, this PR's diff also shows #21088's commits; the commits listed below are this PR's own.

Motivation

Hybrid models with speculative decoding keep per-request replay state next to their recurrent (SSM) state: the inputs of the draft tokens accepted since the last checkpoint, which the verify kernels replay instead of writing a new checkpoint every step. Under KVCacheManagerV2 that state had three problems:

1. It was lost on suspend/resume (Mamba2, GDN). ReplayHistory (Nemotron) and GDNReplayState (Qwen3.5, Qwen4Exp) kept old_x/old_B/old_dt/..., the accepted-token count (PNAT) and the active double-buffer index in standalone tensors indexed by the SSM state slot. A V2 suspend can offload a request's SSM page and resume it into a different slot; the history stayed at the old slot, so the request continued from another request's (or a reset) history.
2. It rode along in every snapshot (KDA). KimiK3 and Glm5Next stored their replay buffers in the REQUIRED SSM slot, so every committed prefix-reuse snapshot also carried replay bytes. Because the replay state still did not follow a request, suspend folded pending drafts into the recurrent state and resume re-seeded the replay caches, so a resumed request matched an unsuspended one only to kernel accuracy.
3. It was sized by the wrong count. The Mamba2/GDN tensors were allocated for every SSM pool slot (snapshot slots included) and counted in no memory budget, although only living requests use replay state.

Design

KVCM: a fixed-capacity PRIVATE SSM pool. KVCacheManagerConfig.max_live_kv_caches (maxLiveKvCaches) bounds the living (active + suspended) KvCaches; createKvCache() rejects one beyond the bound. When it is set, the PRIVATE SSM lifecycle gets a pool group of its own with a fixed GPU capacity of ceil(max_live_kv_caches / max_util_for_resume) slots: PRIVATE pages never enter the radix tree and each living cache holds exactly one, so this keeps the group at or below the resume gate. Construction, adjust()/resize() and the tuner take the fixed group out of the quota first and split the rest by ratio; adjust() never resizes it, and the fixed count is the group's entry in the minimum slots, so a GPU quota too small for all floors is raised, as for the elastic floors. Beam widths above 1 are rejected while a fixed group exists. Without the field, layouts and slot counts are unchanged.

Hybrid managers: replay state as KVCM buffers with its own index.

• MambaHybridCacheManagerV2 gains a replay declaration hook (_get_replay_layout(), _replay_reuse_policy()). Replay roles become buffers of the existing SSM layers, either PRIVATE (own lifecycle, own slot, own index) or coupled (inherit REQUIRED, share the state slot and index).
• The manager resolves the replay group, fills per-request replay slots from get_ssm_block_base_index() next to the state slots, and carries them to the kernels through a static Mamba2Metadata.replay_state_indices buffer (CUDA-graph safe; it aliases state_indices when there is no PRIVATE group) and MambaAcceptanceBatch.replay_indices.
• A request's replay slot is reset the first time it gets one, including generation-only dummies (CUDA-graph padding, warmup, ADP); free and preemption clear that mark, suspend does not.
• max_live_kv_caches is set to the IndexMapper capacity (max_batch_size * pp_size * (2 if disagg or overlap) + reserved dummy slots + guard page), which already bounds every KvCache creation; _create_kv_cache() returns None at the bound like a full IndexMapper. Disaggregated context servers declare no replay buffers.
• Replay bytes are a fixed intercept (ceil(N / util) replay slots) in the manager quota and the KV-cache estimate instead of per-state-slot bytes.
• Python consumers on the hybrid path (pool mapping, page tables, stats, events, disaggregated transfer) resolve groups through get_buffer_group_id / buffer_grouping, re-applying the adoption from [None][feat] Support layer and buffer reuse policies for attention and SSM #21088's history for these sites; connector, sharing and non-hybrid sites keep the legacy API.

Per model family:

Family	Managers	Replay roles	Policy	Kernels taking a replay index
Mamba2	Nemotron	old_x, old_B, old_dt, old_dA_cumsum per layer + 16-byte PNAT/active-buffer row	PRIVATE	replay_selective_state_update (all precompute/main kernels), work-item builders
GDN, Triton replay (default)	Qwen3.5, Qwen4Exp	old_x, old_B, old_dt per layer + 16-byte metadata row (unused old_dA_cumsum dropped)	PRIVATE	cached replay update, all-layer commit (explicit slot/layer strides or descriptor table)
GDN, FlashInfer ring replay (opt-in)	Qwen3.5, Qwen4Exp	k_ring, u_ring, g_ring + metadata row	coupled REQUIRED	none: FlashInfer takes one initial_state_indices; its strided mode reads the KVCM views
KDA	KimiK3, Glm5Next	conv windows, k/g/v/beta caches or draft records, accepted-draft count	PRIVATE	trtllm::kda_mtp_decode*, kda_mtp_commit_pending_drafts, KDA Triton helpers

The KDA suspend-time fold of pending drafts and resume-time re-seed are removed, together with the pending-draft commit path that only they used (kda_mtp_commit_pending_drafts, its replay-only kernel variant and Triton helpers); the disaggregated-transfer seeding is kept. V1 managers are unchanged; every kernel's replay index defaults to the state index.

Commits

Commit	Change
Reserve fixed capacity for private SSM lifecycles	KVCM maxLiveKvCaches, fixed PRIVATE SSM pool group, beam rejection, bindings
Resolve hybrid KV cache groups through buffer APIs	behaviour-neutral consumer adoption for split SSM layers
Keep Mamba2 replay history in private KV cache buffers	generic replay hook, index plumbing, fresh-page reset, max_live_kv_caches, estimator; Mamba2
Keep GDN replay history in KV cache buffers	Qwen3.5 / Qwen4Exp, Triton PRIVATE and FlashInfer ring coupled; experimental engine
Keep KDA replay state in private KV cache buffers	KimiK3 / Glm5Next, CuTe op and Triton helpers
Raise a small GPU quota for the fixed replay pool natively	fixed group joins the native minimum-quota raise; removes a Python workaround
Remove replay leftovers and unrelated SWA scratch check	dead KDA commit path, unused mapping, MambaReplayState protocol, cached group layer ids, shared guard-page parsing; drops an attention-only SWA scratch check that the buffer-API adoption had brought back

Known limitations

• An explicit kv_cache_config.pool_ratio must now include an entry for the PRIVATE replay group (the hot tier ignores it, but the length must match the layer groups).
• The PRIVATE pool is sized by the IndexMapper capacity, which doubles max_batch_size with the overlap scheduler or disaggregation. With prefix reuse enabled this is smaller than before; without it, replay memory can grow up to about 2x.
• FlashInfer ring replay stays coupled to the recurrent-state slot, so its ring bytes are carried by snapshots and come out of the KV budget.
• Disaggregated context-only servers are recognized through TRTLLM_DISAGG_ROLE=context (the existing KDA convention). A context server launched without it declares PRIVATE replay it never uses and bounds living caches by the IndexMapper capacity, so its early index release no longer admits extra requests while transfers are in flight; correctness is unaffected.
• KV connectors are still rejected for hybrid managers (unchanged); the connector layout keeps the layer-level API.

Test Coverage

End-to-end (8x B300, sm_103 wheel): in progress, results will be added here. Planned runs: MTP GSM8K accuracy for Qwen3.5-35B-A3B and Qwen3-Next (Triton and FlashInfer ring replay) and Nemotron-3 Super NVFP4 TP8 with prefix reuse, plus Qwen3.5 and Nemotron runs with a small GPU KV budget and a host tier that force V2 suspend/resume (asserting suspended and resumed requests > 0).

Unit tests (B200, umbriel-b200-043, on the integrated branch):

• C++ kvCacheManagerV2*Test: 208 passed, 0 failed (fixed-capacity sizing, tuner, living-cache bound across suspend, beam rejection, quota raise).
• Python: 4,195 passed, 22 skipped, 92 subtests, 0 failed.
  • KVCM (kv_cache_manager_v2_tests, kvcm2 integration, csa2, estimation, mixed precision, deepseek_v41 snapshots, dual pool, extra buffers): 892 passed, 14 skipped.
  • Mamba managers and kernels (test_mamba_cache_manager.py, test_mamba_model_managers.py, _torch/modules/mamba/): 2,026 passed, 2 skipped.
  • KDA / GLM (_torch/modules/kimi_kda, glm_kpool, Kimi/GLM modeling, test_kda_decode.py): 743 passed, 6 skipped.
  • GDN experimental engine, Qwen4Exp/QSA, disaggregated transfer: 534 passed.
  • One earlier full run hit a single failure of the experimental engine's test_certified_cells[all_layer n=16 T=3 L=8 S=461 ...] (cached-replay kernel); it did not recur in 42 targeted reruns or the final full run.
• New coverage, per family (Mamba2, GDN Triton and ring, KDA):
  • kernels with replay slots disjoint from state slots (strided metadata rows, NaN-filled unused rows), bit-exact against the coupled-index run;
  • suspend, give the request's SSM slot to other caches, resume: replay rows, PNAT and active buffer follow the request, and outputs and states match an unsuspended run bit-exactly;
  • commit for reuse excludes PRIVATE: after the source commits and frees, the PRIVATE pool is all free, reuse_status[private].complete is False and the reader starts from reset bookkeeping;
  • fresh-page reset of generation dummies; max_live_kv_caches equals the IndexMapper capacity and the Python guard returns None.
• FlashInfer 0.6.18 ring replay on the manager's strided KVCM ring views vs contiguous copies: outputs, states and rings bit-exact over 12 steps including flushes.

PR Checklist

• Description explains the behavior and rationale.
• Changes follow the coding guidelines; pre-commit hooks pass.
• Tests cover the new code paths.
• No high-level LLM API contract or API stability reference changes.
• No new dependencies or ownership changes.
• Type stubs and the KVCacheManagerV2 C++ guide are updated.

🤖 Generated with Claude Code