diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index df494fc9f27b..7a720e1e47ee 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -3415,6 +3415,11 @@ def __init__( # update_resources has not recorded yet. self._kda_replay_acceptance_pending: set[int] = set() self._kda_replay_unseeded: set[int] = set() + # Mamba2 / GDN replay state of suspended requests, copied to the host + # by suspend_request in the order of _replay_slot_buffers(); the + # request's next generation step restores it into the slot it resumes + # in. + self._suspended_replay_states: dict[int, list[torch.Tensor]] = {} state_index_capacity = (self.max_batch_size + self._num_reserved_dummy_slots) @@ -4214,6 +4219,19 @@ def _build_cache_config( "are sized once from the SSM pool's slot count, and a KV pool " "rebalance can grow that pool past them: set " "kv_cache_config.enable_kv_pool_rebalance=False.") + if (kv_cache_config.enable_kv_pool_rebalance + and self.local_num_mamba_layers > 0 + and self.replay_step_width is not None): + # Likewise, the Mamba2 / GDN replay tensors (old_x, old_B, old_dt, + # old_dA_cumsum, cache_buf_idx and prev_num_accepted_tokens) are + # allocated once for the SSM pool's slot count, outside the V2 + # pools, and the replay kernels index them by slot without a bound. + raise ValueError( + "The Mamba2 / GDN replay state (speculative decoding with the " + "replay state update) is sized once from the SSM pool's slot " + "count, and a KV pool rebalance can grow that pool past it: " + "set kv_cache_config.enable_kv_pool_rebalance=False, or " + "TRTLLM_USE_MAMBA_REPLAY=0 to use the non-replay path.") cache_tiers = config.cache_tiers gpu_quota = cache_tiers[0].quota minimum_live_quota = self._minimum_live_gpu_quota() @@ -4555,15 +4573,28 @@ def _reset_kda_replay_slots( self.prev_num_accepted_tokens.index_fill_(0, slots, 0) def suspend_request(self, req: LlmRequest) -> None: - """Suspend the request's KV cache with its whole KDA state in V2. - - The drafts the last verify accepted are pending in the slot's KDA - replay buffers, and so is the conv window, which the fused verify does - not write to the V2 conv state. The pending drafts and the conv window - are committed into the slot's SSM and conv states here, before the - pages are released, so the suspended state is self-contained in the - recurrent state; the request's next generation step re-seeds the - replay buffers of whichever slot it resumes in from them. + """Suspend the request's KV cache together with its replay state. + + While the request is suspended, V2 may move its SSM page to another + tier and give its slot to another cache, so the request can resume in + a different slot. Its recurrent state is the slot's V2 SSM and conv + states plus, with speculative decoding, the replay state of the + accepted tokens not yet folded into the SSM state, which must follow + it to the new slot. + + KDA: the pending drafts sit in the slot's KDA replay buffers, and so + does the conv window, which the fused verify does not write to the V2 + conv state. The pending drafts and the conv window are committed into + the slot's SSM and conv states here, before the pages are released, + so the suspended state is self-contained in the recurrent state; the + request's next generation step re-seeds the replay buffers of + whichever slot it resumes in from them. + + Mamba2 / GDN replay: the replay state lives in tensors indexed by slot + outside the V2 pools (``_replay_slot_buffers``). The slot's entries + are copied to the host here, and the request's next generation step + copies them into the slot it resumes in, so the state it resumes with + is bit for bit the one it was suspended with. """ request_id = req.py_request_id kv_cache = self.kv_cache_map.get(request_id) @@ -4585,8 +4616,90 @@ def suspend_request(self, req: LlmRequest) -> None: self._commit_kda_pending_drafts(slot) del self._request_id_to_state_index[request_id] self._kda_replay_suspended.add(request_id) + elif (self._has_replay_slot_buffers and slot >= 0 + and kv_cache is not None and kv_cache.is_active + and not req.is_dummy): + # Forgetting the slot also keeps a later suspend before the next + # generation step from overwriting the copy with the new slot's + # entries. + self._suspended_replay_states[request_id] = ( + self._copy_replay_slot_to_host(slot)) + del self._request_id_to_state_index[request_id] super().suspend_request(req) + @property + def _has_replay_slot_buffers(self) -> bool: + """Whether this rank keeps Mamba2 / GDN replay state in tensors + indexed by state slot outside the V2 pools.""" + return (not self.use_kda_replay_update + and self.prev_num_accepted_tokens is not None) + + def _replay_slot_buffers(self) -> list[tuple[torch.Tensor, int]]: + """The Mamba2 / GDN replay tensors, each with its slot dimension. + + ``prev_num_accepted_tokens`` and ``cache_buf_idx`` are ``[slots]``; + ``old_x``, ``old_B``, ``old_dt`` and ``old_dA_cumsum`` are + ``[layers, slots, 2, ...]``, both halves of the double-buffered + history. + """ + buffers = [(self.prev_num_accepted_tokens, 0), (self.cache_buf_idx, 0)] + buffers.extend((history, 1) + for history in (self.old_x, self.old_B, self.old_dt, + self.old_dA_cumsum) + if history is not None) + return buffers + + def _copy_replay_slot_to_host(self, slot: int) -> list[torch.Tensor]: + """Copy the slot's entries of ``_replay_slot_buffers()`` to new host + tensors, in that order. + + Only the current stream (verify, context reset) and the KV caches' + stream (``_restore_suspended_replay_states``) write these tensors. The + copies run on the KV caches' stream after the current stream's earlier + work, which includes the request's last verify, and before its later + work and the KV caches' stream's later work, which include whatever + the slot's next owner writes to it. + """ + host_entries = [] + with self._on_kv_cache_stream(): + for buffer, slot_dim in self._replay_slot_buffers(): + entry = buffer.select(slot_dim, slot) + host_entry = torch.empty(entry.shape, + dtype=entry.dtype, + pin_memory=prefer_pinned()) + host_entry.copy_(entry, non_blocking=True) + host_entries.append(host_entry) + return host_entries + + def _restore_suspended_replay_states(self, requests: list[LlmRequest], + num_contexts: int) -> None: + """Copy the replay state saved by ``suspend_request`` into the slots + the batch's suspended generation requests resume in. + + A suspended request that comes back as a context request starts over, + so its saved state is dropped. The copies run on the KV caches' + stream, after the copies out of the slots' previous owners and the + current stream's earlier work, and before the current stream's later + work, which includes the next verify. + """ + if not self._suspended_replay_states: + return + for request in requests[:num_contexts]: + self._suspended_replay_states.pop(request.py_request_id, None) + resumed = [(self._suspended_replay_states.pop(request.py_request_id), + self._request_id_to_state_index[request.py_request_id]) + for request in requests[num_contexts:] + if request.py_request_id in self._suspended_replay_states] + if not resumed: + return + with self._on_kv_cache_stream(): + for host_entries, slot in resumed: + for (buffer, + slot_dim), host_entry in zip(self._replay_slot_buffers(), + host_entries): + buffer.select(slot_dim, slot).copy_(host_entry, + non_blocking=True) + @contextmanager def _on_kv_cache_stream(self): """Run the enclosed work on the KV caches' stream, after the current @@ -4682,6 +4795,7 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): self._kda_replay_suspended.discard(request.py_request_id) self._kda_replay_acceptance_pending.discard(request.py_request_id) self._kda_replay_unseeded.discard(request.py_request_id) + self._suspended_replay_states.pop(request.py_request_id, None) super().free_resources(request, pin_on_release) def prepare_resources(self, scheduled_batch: ScheduledRequests): @@ -4748,6 +4862,9 @@ def _setup_state_indices(self, self._request_id_to_state_index[req.py_request_id] = value self._request_id_to_is_dummy[req.py_request_id] = dummy + if self._has_replay_slot_buffers: + self._restore_suspended_replay_states(requests, num_contexts) + if self.use_kda_replay_update: # A suspended request that comes back as a context request starts # over. One that comes back generating has its whole KDA state in diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 9c8067141c5e..3ca8d42b5604 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -3902,6 +3902,26 @@ def test_v2_kda_replay_refuses_kv_pool_rebalance(kda_replay): ) +@skip_no_cuda +@pytest.mark.parametrize("conv_state_layout", ["x_b_c", "q_k_v"], ids=["mamba2", "gdn"]) +@pytest.mark.parametrize("use_replay_state_update", [True, False], ids=["replay", "no_replay"]) +def test_v2_replay_refuses_kv_pool_rebalance(conv_state_layout, use_replay_state_update): + """The Mamba2 / GDN replay state holds one entry per SSM slot the pool keeps at construction, outside the V2 + pools, and a KV pool rebalance can grow that pool past it, so the manager refuses enable_kv_pool_rebalance with + it; speculative decoding without it keeps the rebalance allowed.""" + kwargs = dict( + spec_config=MTPDecodingConfig(max_draft_len=2), + use_replay_state_update=use_replay_state_update, + conv_state_layout=conv_state_layout, + enable_kv_pool_rebalance=True, + ) + if not use_replay_state_update: + _build_v2_hybrid_with_mamba_layer(**kwargs).shutdown() + return + with pytest.raises(ValueError, match="enable_kv_pool_rebalance=False"): + _build_v2_hybrid_with_mamba_layer(**kwargs) + + def _v2_kda_draft_record_manager(kda_draft_records, enable_block_reuse): """A V2 hybrid manager stub for _build_cache_config: one KDA layer with the replay caches and one attention layer.""" @@ -4483,6 +4503,195 @@ def test_v2_kda_replay_suspend_resume_matches_an_unsuspended_run(slot_taken): torch.testing.assert_close(states, expected_states, atol=1e-3, rtol=1e-3) +_REPLAY_HISTORIES = ("old_x", "old_B", "old_dt", "old_dA_cumsum") +_REPLAY_LAYOUTS = pytest.mark.parametrize( + "conv_state_layout", ["x_b_c", "q_k_v"], ids=["mamba2_replay", "gdn_replay"] +) + + +def _v2_replay_manager_with_large_slots(conv_state_layout): + """_v2_kda_replay_manager_with_large_slots' geometry with the Mamba2 (``x_b_c``) or GDN (``q_k_v``) replay state + update in place of the KDA replay caches: an SSM slot (~2.4 MiB) is larger than the 2 MiB allocation grain, so the + SSM pool runs out of slots before the attention pages run out.""" + num_mamba_layers = 6 + return MambaHybridCacheManagerV2( + mamba_d_state=128, + mamba_d_conv=4, + mamba_num_heads=6, + mamba_n_groups=6, + mamba_head_dim=128, + mamba_num_layers=num_mamba_layers, + mamba_layer_mask=[True] * num_mamba_layers + [False], + mamba_cache_dtype=torch.bfloat16, + mamba_ssm_cache_dtype=torch.float32, + kv_cache_config=KvCacheConfig( + max_tokens=4096, enable_block_reuse=False, max_util_for_resume=1.0 + ), + kv_cache_type=CacheTypeCpp.SELF, + num_layers=1, + num_kv_heads=4, + head_dim=64, + tokens_per_block=32, + max_seq_len=128, + max_batch_size=4, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.HALF, + spec_config=MTPDecodingConfig(max_draft_len=2), + layer_mask=[False] * num_mamba_layers + [True], + vocab_size=1024, + conv_state_layout=conv_state_layout, + use_replay_state_update=True, + ) + + +def _add_generation_request(mgr, request_id): + """A generation request holding a slot. add_dummy_requests marks its requests as dummies, which suspend_request + leaves alone, so the request is turned into a real one.""" + (request,) = mgr.add_dummy_requests( + [request_id], token_nums=[8], is_gen=True, max_num_draft_tokens=2 + ) + request.is_dummy_request = False + mgr._setup_state_indices([request], num_contexts=0) + return request + + +def _fill_replay_slot(mgr, slot, generator, num_accepted, cache_buf_idx): + """Random V2 SSM and conv states and replay histories at ``slot``, with ``num_accepted`` accepted tokens pending + in history buffer ``cache_buf_idx``.""" + for ssm_states, conv_states in zip(mgr.all_ssm_states, mgr.all_conv_states): + ssm_states[slot] = torch.randn(ssm_states.shape[1:], generator=generator, device="cuda") + conv_states[slot] = torch.randn( + conv_states.shape[1:], generator=generator, device="cuda" + ).to(conv_states.dtype) + mgr.prev_num_accepted_tokens[slot] = num_accepted + mgr.cache_buf_idx[slot] = cache_buf_idx + for name in _REPLAY_HISTORIES: + history = getattr(mgr, name) + history[:, slot] = torch.randn( + history[:, slot].shape, generator=generator, device="cuda" + ).to(history.dtype) + + +def _replay_slot_state(mgr, slot): + """The whole recurrent state at ``slot``: the V2 SSM and conv states and the replay state.""" + torch.cuda.synchronize() + return { + "ssm": torch.stack([ssm_states[slot] for ssm_states in mgr.all_ssm_states]).clone(), + "conv": torch.stack([conv_states[slot] for conv_states in mgr.all_conv_states]).clone(), + "prev_num_accepted_tokens": mgr.prev_num_accepted_tokens[slot].clone(), + "cache_buf_idx": mgr.cache_buf_idx[slot].clone(), + **{name: getattr(mgr, name)[:, slot].clone() for name in _REPLAY_HISTORIES}, + } + + +def _assert_same_replay_slot_state(actual, expected): + for name, value in expected.items(): + assert torch.equal(actual[name], value), name + + +@skip_no_cuda +@_REPLAY_LAYOUTS +@pytest.mark.parametrize("slot_taken", [True, False], ids=["slot_taken", "slot_kept"]) +def test_v2_replay_suspend_resume_keeps_the_replay_state(conv_state_layout, slot_taken): + """A generation request is suspended with accepted tokens pending in its replay history and resumes. Its V2 SSM + and conv states come back with its SSM page, and its replay state (prev_num_accepted_tokens, cache_buf_idx and both + halves of every replay history), which lives outside the V2 pools, comes back with them, bit for bit, in the slot + it resumes in. + + slot_taken: while the request is suspended, new caches fill the SSM pool, so the V2 runtime moves its SSM page to + the host tier and gives its slot to one of them, whose prefill and verify rounds overwrite the slot; the request + resumes in another slot, and the restore leaves the new owner's state alone. slot_kept: nothing else runs while it + is suspended.""" + mgr = _v2_replay_manager_with_large_slots(conv_state_layout) + others = [] + request = None + try: + request = _add_generation_request(mgr, 1) + generator = torch.Generator(device="cuda").manual_seed(0) + slot = mgr._request_id_to_state_index[1] + _fill_replay_slot(mgr, slot, generator, num_accepted=5, cache_buf_idx=1) + expected = _replay_slot_state(mgr, slot) + + mgr.suspend_request(request) + if slot_taken: + _take_ssm_slot(mgr, slot, others) + _fill_replay_slot(mgr, slot, generator, num_accepted=3, cache_buf_idx=0) + new_owner_state = _replay_slot_state(mgr, slot) + assert mgr.resume_request(request) + mgr._setup_state_indices([request], num_contexts=0) + + resumed_slot = mgr._request_id_to_state_index[1] + if slot_taken: + assert resumed_slot != slot + _assert_same_replay_slot_state(_replay_slot_state(mgr, slot), new_owner_state) + _assert_same_replay_slot_state(_replay_slot_state(mgr, resumed_slot), expected) + finally: + for cache in others: + cache.close() + if request is not None: + mgr.free_resources(request) + mgr.shutdown() + + +@skip_no_cuda +@_REPLAY_LAYOUTS +def test_v2_replay_suspended_request_resumed_as_context_starts_over(conv_state_layout): + """A suspended request that comes back as a context request re-prefills: its slot's replay state is reset, and + the replay state saved at the suspend is dropped instead of restored, now or at its next generation step.""" + mgr = _v2_replay_manager_with_large_slots(conv_state_layout) + request = None + try: + request = _add_generation_request(mgr, 1) + generator = torch.Generator(device="cuda").manual_seed(0) + _fill_replay_slot( + mgr, mgr._request_id_to_state_index[1], generator, num_accepted=5, cache_buf_idx=1 + ) + mgr.suspend_request(request) + assert 1 not in mgr._request_id_to_state_index + assert mgr.resume_request(request) + + mgr._setup_state_indices([request], num_contexts=1) + mgr._reset_context_mamba_slots(1) + assert not mgr._suspended_replay_states + mgr._setup_state_indices([request], num_contexts=0) + + state = _replay_slot_state(mgr, mgr._request_id_to_state_index[1]) + for name in ("prev_num_accepted_tokens", "cache_buf_idx", *_REPLAY_HISTORIES): + assert not state[name].any(), name + finally: + if request is not None: + mgr.free_resources(request) + mgr.shutdown() + + +@skip_no_cuda +@_REPLAY_LAYOUTS +def test_v2_replay_suspend_keeps_nothing_for_freed_or_dummy_requests(conv_state_layout): + """The replay state saved at a suspend is released with the request when it is freed while suspended. A dummy + request's replay state is never read, so its suspend saves nothing and keeps its slot mapping.""" + mgr = _v2_replay_manager_with_large_slots(conv_state_layout) + request = None + dummy = None + try: + request = _add_generation_request(mgr, 1) + mgr.suspend_request(request) + assert set(mgr._suspended_replay_states) == {1} + mgr.free_resources(request) + request = None + assert not mgr._suspended_replay_states + + (dummy,) = mgr.add_dummy_requests([2], token_nums=[8], is_gen=True, max_num_draft_tokens=2) + slot = mgr._request_id_to_state_index[2] + mgr.suspend_request(dummy) + assert not mgr._suspended_replay_states + assert mgr._request_id_to_state_index[2] == slot + finally: + for live in (request, dummy): + if live is not None: + mgr.free_resources(live) + mgr.shutdown() + + @skip_no_cuda def test_v2_kda_replay_commit_without_the_replay_caches(): """With the verify kernels' draft records (``kda_state_tok``) in place of the per-draft replay caches, the