From 1da7f0f75725d9bd481e91ad9bfb5bdaa63d633e Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Thu, 29 Oct 2026 12:00:04 -0700 Subject: [PATCH 01/14] [TRTLLM-17390][fix] Size the separate draft-KV mirror from request history A one-engine external drafter's mirror follows the target's position and chunk, and context resizes pass the chunk start as history so sliding-window draft pools hold window plus chunk; a disaggregated generation-init prompt counts as history, and dummy draft mirrors take the target's history hint. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- .../kv_cache/kv_cache_manager_v2.py | 87 +++-- .../kv_cache/test_kvcm2_integration.py | 297 ++++++++++++++++++ .../test_draft_kv_resume_pressure.py | 36 +++ 3 files changed, 402 insertions(+), 18 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index a4f91d5b1e8f..a51a8998e3fc 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -2053,6 +2053,14 @@ def __init__( from tensorrt_llm._torch.speculative import draft_prompt_lookahead, get_num_extra_kv_tokens self.num_extra_kv_tokens = get_num_extra_kv_tokens(spec_config) + # The draft request view (LlmRequest.use_draft_model) reads the context + # position and chunk a separate draft engine advances for its own + # prefill; a one-engine external drafter (DFlash, DSpark) has no such + # engine, so those fields stay at position 0 / chunk = prompt length + # and its mirror follows the target's chunking instead. + self._mirror_uses_draft_view = not ( + is_draft and spec_config is not None and spec_config.spec_dec_mode.is_external_drafter() + ) # Mirror V1: expose max_draft_len so the native disagg AuxBuffer # (_make_aux_buffer's getattr fallback) is sized for MTP/spec decoding. self.max_draft_len = spec_config.max_draft_len if spec_config is not None else 0 @@ -6059,16 +6067,22 @@ def _draft_pool_diagnostic(self) -> str: byte budget, and the capacity scheduler admits on the TARGET pool alone (`scheduler_v2` touches `draft_kv_cache_manager` only to suspend/free). A draft pool smaller in tokens than the target cannot backpressure -- it can - only raise, and the raise kills every rank. + only raise (see _mirror_resize_error), and the raise kills every rank. The first question is always how big the draft pool was and how full, so the message answers it rather than leaving post-hoc arithmetic over the split log. """ - live = sum(c.capacity for c in self.kv_cache_map.values()) + capacity = sum(c.capacity for c in self.kv_cache_map.values()) + blocks = sum(int(c.num_blocks) for c in self.kv_cache_map.values()) + per_cache = ", ".join( + f"{req_id}: capacity={c.capacity} history={c.history_length} blocks={c.num_blocks}" + for req_id, c in self.kv_cache_map.items() + ) return ( - f" [draft pool: {len(self.kv_cache_map)} live caches holding " - f"{live} tokens, gpu_max_tokens={self._gpu_max_tokens}]" + f" [draft pool: {len(self.kv_cache_map)} live caches, capacity {capacity} tokens " + f"in {blocks} allocated blocks ({per_cache}), " + f"windows={self.max_attention_window_vec}, gpu_max_tokens={self._gpu_max_tokens}]" ) def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): @@ -6080,8 +6094,11 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): scheduler pre-creates context caches to coordinate the reusable prefix. """ # Joint V2 reuse deliberately shares the target request progress. A - # separate draft engine still uses the legacy draft request view. - with request_context(not self.enable_joint_kv_cache_reuse, scheduled_batch): + # separate draft engine still uses the legacy draft request view; a + # one-engine external drafter never advances that view, so its mirror + # reads the target's position and chunk (see _mirror_uses_draft_view). + use_draft_view = self._mirror_uses_draft_view and not self.enable_joint_kv_cache_reuse + with request_context(use_draft_view, scheduled_batch): for req in scheduled_batch.context_requests: kv_cache = self._mirror_draft_kv_cache(req) if kv_cache is None: @@ -6114,12 +6131,27 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): + draft_len + self.num_extra_kv_tokens ) - if not kv_cache.resize(capacity): - raise RuntimeError( - f"Draft KV cache context resize failed for request " - f"{req.py_request_id}: could not resize to {capacity} tokens" - f"{self._draft_pool_diagnostic()}" - ) + # The chunk start is this cache's history: sliding-window layers + # keep pages only from history + 1 - window (AttnLifeCycle:: + # getStaleRange), so a mirror holds its window plus the chunk + # it is computing rather than the whole prompt so far. + # The executor advances the draft history through + # update_context_resources only under joint reuse, and a reused + # target prefix starts the mirror at position > 0 with no pages + # of its own, so the resize carries the history itself. + history = max(kv_cache.history_length, req.context_current_position) + if req.is_disagg_generation_init_state: + # A disaggregated generation worker receives the prompt KV + # instead of prefilling it, and the standalone drafter + # starts this request on a fresh slot (DFlash/DSpark + # prepare): the mirror never holds drafter K/V for the + # transferred prompt. The whole prompt is therefore + # history, so the window trims the mirror to its tail + # plus the generation headroom instead of one prompt-sized + # chunk that a windowed draft pool cannot hold. + history = max(history, req.context_current_position + req.context_chunk_size) + if not kv_cache.resize(capacity, history): + raise self._mirror_resize_error(req, "context", capacity) self._after_shrink(req.py_request_id, kv_cache) for req in scheduled_batch.generation_requests: @@ -6143,11 +6175,25 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): continue new_cap = self._required_gen_capacity(req, kv_cache.capacity) if not kv_cache.resize(new_cap): - raise RuntimeError( - f"Draft KV cache generation resize failed for request " - f"{req.py_request_id}: could not resize to {new_cap} tokens" - f"{self._draft_pool_diagnostic()}" - ) + raise self._mirror_resize_error(req, "generation", new_cap) + + def _mirror_resize_error(self, req: LlmRequest, phase: str, capacity: int) -> RuntimeError: + """The error for a draft mirror that could not grow to ``capacity``. + + `KvCache.resize` leaves the cache as it was when the pool has no page to + give, and the cache as it was cannot hold what the drafter writes next: + the DFlash block decode appends the block's own K/V past the stored + context through the mirror's block table, so a mirror that did not grow + (a fresh one has no page at all) hands that write a table entry with no + page behind it (``BAD_PAGE_INDEX`` on the python V2 backend, another + request's block 0 on the cpp one). There is no per-request way to skip + drafting, so the failure stays explicit. The message carries the pool + diagnostic. + """ + return RuntimeError( + f"Draft KV cache {phase} resize failed for request {req.py_request_id}: " + f"could not resize to {capacity} tokens{self._draft_pool_diagnostic()}" + ) def _reuse_token_source(self, req: LlmRequest) -> Sequence[int]: """Return a zero-copy beam-0 token view for block reuse.""" @@ -7015,7 +7061,12 @@ def release_resources( if not success: release_resources(req, free_draft_resources=True) return None - success = draft_kv_cache.resize(dummy_capacity) + if is_gen and not materialize_history: + draft_kv_cache.enable_swa_scratch_reuse = False + # Same history hint as the target cache above: a sliding-window + # draft pool then backs only the dummy's window, not its whole + # length. + success = draft_kv_cache.resize(dummy_capacity, history_length=history_hint) if not success: release_resources(req, free_draft_resources=True) return None diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py index 5394f95c732f..64afc689c57a 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py @@ -33,6 +33,7 @@ KVCacheManagerV2, Role, _extend_swa_windows_for_reuse, + _get_generation_kv_capacity, _KVCacheManagerInitStatus, _swa_branch_snapshot_point, _swa_endpoint_priority, @@ -1885,6 +1886,7 @@ def test_draft_manager_keeps_shared_progress_across_context_and_generation() -> manager.is_draft = True manager.enable_block_reuse = True manager.enable_joint_kv_cache_reuse = True + manager._mirror_uses_draft_view = True manager.kv_cache_map = {request.py_request_id: kv_cache} observed_progress_views = [] @@ -1910,6 +1912,301 @@ def resume_and_restore(request_id, current_cache): manager._required_gen_capacity.assert_not_called() +@pytest.mark.parametrize( + ("mirror_history", "expected_history"), + [ + # A fresh mirror for a request whose target prefix was reused: the + # history jumps to the chunk start so no page below the window is + # allocated for [0, position). + (0, 8192), + # A later chunk of the same request: the history follows the chunk. + (4096, 8192), + # Joint reuse already advanced the history past the chunk start; a + # resize may not decrease it. + (8704, 8704), + ], +) +def test_unpaired_one_engine_draft_mirror_follows_the_target_chunk( + mirror_history: int, expected_history: int +) -> None: + """A one-engine external drafter's mirror reads the target's position and + chunk, and its context resize carries the chunk start as history.""" + position, chunk, draft_len = 8192, 1024, 3 + request = LlmRequest( + request_id=41, + max_new_tokens=4, + input_tokens=[0] * (position + chunk), + sampling_config=SamplingConfig(1), + is_streaming=False, + draft_tokens=[1] * draft_len, + ) + request.state = LlmRequestState.CONTEXT_INIT + request.context_current_position = position + request.context_chunk_size = chunk + assert request.is_last_context_chunk + + kv_cache = Mock(is_active=True, capacity=mirror_history, history_length=mirror_history) + kv_cache.resize.return_value = True + manager = object.__new__(KVCacheManagerV2) + manager.is_draft = True + manager.enable_joint_kv_cache_reuse = False + manager._mirror_uses_draft_view = False + manager.num_extra_kv_tokens = 0 + manager.kv_cache_map = {request.py_request_id: kv_cache} + manager._resume_and_restore = Mock(return_value=True) + + batch = ScheduledRequests() + batch.context_requests_last_chunk.append(request) + manager._prepare_draft_resources(batch) + + kv_cache.resize.assert_called_once_with(position + chunk + draft_len, expected_history) + assert not request.use_draft_model + + +def test_unpaired_two_engine_draft_mirror_reads_the_draft_request_view() -> None: + """The draft view starts at position 0 with the whole prompt as its chunk + until a separate draft engine advances it, so a two-engine mirror sizes + its cache for the prompt at once.""" + position, chunk, draft_len = 8192, 1024, 3 + request = LlmRequest( + request_id=42, + max_new_tokens=4, + input_tokens=[0] * (position + chunk), + sampling_config=SamplingConfig(1), + is_streaming=False, + draft_tokens=[1] * draft_len, + ) + request.state = LlmRequestState.CONTEXT_INIT + request.context_current_position = position + request.context_chunk_size = chunk + + kv_cache = Mock(is_active=True, capacity=0, history_length=0) + kv_cache.resize.return_value = True + manager = object.__new__(KVCacheManagerV2) + manager.is_draft = True + manager.enable_joint_kv_cache_reuse = False + manager._mirror_uses_draft_view = True + manager.num_extra_kv_tokens = 0 + manager.kv_cache_map = {request.py_request_id: kv_cache} + manager._resume_and_restore = Mock(return_value=True) + + batch = ScheduledRequests() + batch.context_requests_last_chunk.append(request) + manager._prepare_draft_resources(batch) + + kv_cache.resize.assert_called_once_with(position + chunk + draft_len, 0) + assert not request.use_draft_model + + +@pytest.mark.parametrize("phase", ["context", "generation"]) +def test_unpaired_draft_mirror_that_cannot_grow_raises_with_the_pool_diagnostic(phase: str) -> None: + """When a mirror finds no page to grow into, the iteration fails + explicitly: the mirror left at its size has no storage for the K/V the + drafter writes next, and there is no per-request way to skip drafting. + The error carries the pool diagnostic.""" + position, chunk, draft_len = 8192, 1024, 3 + + def make_request(request_id: int) -> LlmRequest: + request = LlmRequest( + request_id=request_id, + max_new_tokens=4, + input_tokens=[0] * (position + chunk), + sampling_config=SamplingConfig(1), + is_streaming=False, + draft_tokens=[1] * draft_len, + ) + if phase == "context": + request.state = LlmRequestState.CONTEXT_INIT + request.context_current_position = position + request.context_chunk_size = chunk + else: + request.state = LlmRequestState.GENERATION_IN_PROGRESS + return request + + requests = [make_request(51), make_request(52)] + caches = { + 51: Mock(is_active=True, capacity=position, history_length=position, beam_width=1), + 52: Mock(is_active=True, capacity=position, history_length=position, beam_width=1), + } + caches[51].resize.return_value = True + caches[52].resize.return_value = False + manager = object.__new__(KVCacheManagerV2) + manager.is_draft = True + manager.enable_joint_kv_cache_reuse = False + manager._mirror_uses_draft_view = False + manager.num_extra_kv_tokens = 0 + manager.kv_cache_map = dict(caches) + manager._allocated_draft_lens = set() + manager._resume_and_restore = Mock(return_value=True) + manager._required_gen_capacity = Mock(return_value=position + 1 + draft_len) + manager._draft_pool_diagnostic = Mock(return_value=" [draft pool: full]") + manager.kv_cache_type = CacheType.SELF + manager.max_beam_width = 1 + + batch = ScheduledRequests() + if phase == "context": + batch.context_requests_last_chunk.extend(requests) + expected_resize = call(position + chunk + draft_len, position) + expected_capacity = position + chunk + draft_len + else: + batch.generation_requests.extend(requests) + expected_resize = call(position + 1 + draft_len) + expected_capacity = position + 1 + draft_len + with pytest.raises(RuntimeError) as excinfo: + manager._prepare_draft_resources(batch) + + assert caches[51].resize.call_args_list == [expected_resize] + assert caches[52].resize.call_args_list == [expected_resize] + message = str(excinfo.value) + assert f"Draft KV cache {phase} resize failed for request 52" in message + assert f"could not resize to {expected_capacity} tokens" in message + assert "[draft pool: full]" in message + + +def test_exhausted_pool_fails_explicitly_on_an_empty_fresh_mirror() -> None: + """A fresh mirror whose first resize fails holds no page at all. The + manager raises before the forward, with a diagnostic that shows the empty + mirror and how full the pool was.""" + prompt, draft_len = 1024, 3 + request = LlmRequest( + request_id=61, + max_new_tokens=4, + input_tokens=[0] * prompt, + sampling_config=SamplingConfig(1), + is_streaming=False, + draft_tokens=[1] * draft_len, + ) + request.state = LlmRequestState.CONTEXT_INIT + request.context_current_position = 0 + request.context_chunk_size = prompt + + occupant = Mock(is_active=True, capacity=1216, history_length=0, num_blocks=38) + fresh = Mock(is_active=True, capacity=0, history_length=0, num_blocks=0) + fresh.resize.return_value = False + manager = object.__new__(KVCacheManagerV2) + manager.is_draft = True + manager.enable_joint_kv_cache_reuse = False + manager._mirror_uses_draft_view = False + manager.num_extra_kv_tokens = 0 + manager.kv_cache_map = {60: occupant, 61: fresh} + manager.max_attention_window_vec = [2048] + manager._gpu_max_tokens = 1216 + manager._resume_and_restore = Mock(return_value=True) + + batch = ScheduledRequests() + batch.context_requests_last_chunk.append(request) + with pytest.raises(RuntimeError) as excinfo: + manager._prepare_draft_resources(batch) + + fresh.resize.assert_called_once_with(prompt + draft_len, 0) + message = str(excinfo.value) + assert ( + f"context resize failed for request 61: could not resize to {prompt + draft_len} tokens" + in message + ) + assert "2 live caches, capacity 1216 tokens in 38 allocated blocks" in message + assert "61: capacity=0 history=0 blocks=0" in message + assert "windows=[2048]" in message + assert "gpu_max_tokens=1216" in message + + +def test_windowed_draft_pool_backs_only_the_window_of_a_generation_dummy() -> None: + """A CUDA-graph warmup dummy is a generation request at the longest length + the target pool can hold. Its draft mirror carries the same history hint + as the target cache, so a windowed draft pool sized for a few mirrors of + window + chunk backs only the dummy's window; without the hint the mirror + would need the dummy's whole length, which such a pool cannot hold.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + spec = DFlashDecodingConfig(max_draft_len=4, speculative_model="draft") + # 32 KiB per token (8 layers, 8 KV heads of 128, fp16): a mirror of a + # 256-token window plus a 512-token chunk is ~26 MiB; the pool holds four + # of them (~104 MiB), while a 60k-token dummy held in full needs ~1.8 GiB. + num_layers, num_kv_heads, head_dim = 8, 8, 128 + tokens_per_block, window, chunk, max_seq_len, batch_size = 32, 256, 512, 65536, 2 + token_num = 60_000 + draft_config = KvCacheConfig( + enable_block_reuse=False, max_attention_window=[window] * num_layers + ) + bytes_per_token = num_layers * 2 * num_kv_heads * head_dim * 2 + _, headroom = _get_generation_kv_capacity(spec, is_draft=True) + tokens_per_mirror = window + chunk + 2 * headroom + tokens_per_block + pages_per_mirror = -(-tokens_per_mirror // tokens_per_block) + draft_budget = 2 * batch_size * pages_per_mirror * tokens_per_block * bytes_per_token + assert token_num * bytes_per_token > 4 * draft_budget + draft = KVCacheManagerV2( + draft_config.model_copy(update={"max_gpu_total_bytes": draft_budget}), + CacheType.SELF, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=batch_size, + max_num_tokens=chunk, + mapping=Mapping(), + dtype=DataType.HALF, + spec_config=spec, + is_draft=True, + ) + # A small windowless target pool (512 B per token) that holds the whole + # dummy, so the target side of add_dummy_requests cannot be what fails. + target = KVCacheManagerV2( + KvCacheConfig(enable_block_reuse=False, max_gpu_total_bytes=128 << 20), + CacheType.SELF, + num_layers=1, + num_kv_heads=1, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=batch_size, + max_num_tokens=chunk, + mapping=Mapping(), + dtype=DataType.HALF, + spec_config=spec, + ) + # Only the external drafter's own manager follows the target's chunking. + assert not draft._mirror_uses_draft_view + assert target._mirror_uses_draft_view + stream = torch.cuda.current_stream().cuda_stream + + def live_pages(cache: object) -> int: + return len(cache.get_aggregated_page_indices(0, valid_only=True)) + + requests = None + try: + # The mirror resize without a history hint: the whole dummy on the + # windowed pool, which is more than the pool has. + probe = draft.impl.create_kv_cache() + try: + assert probe.resume(stream) + assert not probe.resize(token_num + draft.num_extra_kv_tokens) + finally: + probe.close() + requests = target.add_dummy_requests( + [7], + token_nums=[token_num], + is_gen=True, + max_num_draft_tokens=spec.max_draft_len, + draft_kv_cache_manager=draft, + ) + assert requests is not None, "the dummy's draft mirror did not fit the windowed pool" + (request,) = requests + target_cache = target.kv_cache_map[request.py_request_id] + mirror = draft.kv_cache_map[request.py_request_id] + assert live_pages(target_cache) >= token_num // tokens_per_block + assert mirror.capacity == target_cache.capacity + assert mirror.history_length == token_num - 1 + assert live_pages(mirror) <= pages_per_mirror + finally: + for request in requests or (): + draft.free_resources(request) + target.free_resources(request) + draft.shutdown() + target.shutdown() + + def _make_publishing_manager(policy: BlockReusePolicy) -> KVCacheManagerV2: """A manager wired just far enough to run the publish/history bookkeeping.""" manager = object.__new__(KVCacheManagerV2) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_kv_resume_pressure.py b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_kv_resume_pressure.py index dc90385bf0cb..2daeb957320d 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_kv_resume_pressure.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_kv_resume_pressure.py @@ -54,6 +54,7 @@ def __init__(self, req_id): self.context_current_position = 0 self.context_chunk_size = 8 self.py_draft_tokens = None + self.is_disagg_generation_init_state = False class _Batch: @@ -78,10 +79,12 @@ def _manager(resume_ok): manager.is_draft = True manager.enable_joint_kv_cache_reuse = False manager.num_extra_kv_tokens = 0 + manager._mirror_uses_draft_view = False # one-engine external drafter: target view manager._allocated_draft_lens = set() cache = MagicMock(name="kv_cache") cache.resize.return_value = True cache.capacity = 0 + cache.history_length = 0 manager._mirror_draft_kv_cache = MagicMock(return_value=cache) manager._resume_and_restore = MagicMock(return_value=resume_ok) manager._required_gen_capacity = MagicMock(return_value=1) @@ -125,6 +128,39 @@ def test_context_request_is_prepared_normally_when_resume_succeeds(): cache.resize.assert_called_once() +def test_disagg_generation_init_sizes_the_mirror_from_history_and_chunk() -> None: + """A transferred prompt is history for the drafter's mirror, not a chunk to hold. + + On a disaggregated generation worker the prompt KV arrives over the wire + and the one-engine drafter starts the request on a fresh slot, so the + mirror never holds drafter K/V for that prompt. The resize therefore + carries the whole prompt as history, which lets the attention window trim + the mirror to its tail plus the generation headroom; a history of 0 asks + a windowed draft pool for one prompt-sized chunk instead. + """ + manager, cache = _manager(resume_ok=True) + req = _Req(9) + req.is_disagg_generation_init_state = True + req.context_current_position = 0 + req.context_chunk_size = 200_000 + + manager._prepare_draft_resources(_Batch(context_requests=[req])) + + cache.resize.assert_called_once_with(200_000, 200_000) + + +def test_ordinary_context_chunk_keeps_its_start_as_history() -> None: + """The aggregated chunked path: history is the chunk start, not its end.""" + manager, cache = _manager(resume_ok=True) + req = _Req(10) + req.context_current_position = 64 + req.context_chunk_size = 8 + + manager._prepare_draft_resources(_Batch(context_requests=[req])) + + cache.resize.assert_called_once_with(72, 64) + + def test_generation_request_still_raises_on_a_resume_refusal(): """Skipping a generation request asserts later in C++, so it must raise here.""" manager, _ = _manager(resume_ok=False) From 2ce0ee91cd8270d6392e1c27281acb883d01db10 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Thu, 29 Oct 2026 13:47:05 -0700 Subject: [PATCH 02/14] [TRTLLM-17390][fix] Size draft mirrors only where the drafter addresses them Context history applies only to external drafters' mirrors that already hold [0, chunk start); MTP/Eagle keep their resize. A disaggregated generation-init mirror is sized in the drafter's slot-local positions instead of trimming the pages that slot writes. Tests pin both on a real windowed pool and drafter. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- .../kv_cache/kv_cache_manager_v2.py | 62 ++- .../kv_cache/test_kvcm2_integration.py | 495 +++++++++++++++--- .../test_draft_kv_resume_pressure.py | 37 +- 3 files changed, 465 insertions(+), 129 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index a51a8998e3fc..4cf63a687a8d 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -2055,9 +2055,10 @@ def __init__( self.num_extra_kv_tokens = get_num_extra_kv_tokens(spec_config) # The draft request view (LlmRequest.use_draft_model) reads the context # position and chunk a separate draft engine advances for its own - # prefill; a one-engine external drafter (DFlash, DSpark) has no such - # engine, so those fields stay at position 0 / chunk = prompt length - # and its mirror follows the target's chunking instead. + # prefill; a one-engine external drafter (DFlash, DSpark, PARD, + # DraftTarget one-model) has no such engine, so those fields stay at + # position 0 / chunk = prompt length and its mirror follows the + # target's chunking instead. self._mirror_uses_draft_view = not ( is_draft and spec_config is not None and spec_config.spec_dec_mode.is_external_drafter() ) @@ -6073,15 +6074,19 @@ def _draft_pool_diagnostic(self) -> str: message answers it rather than leaving post-hoc arithmetic over the split log. """ + # num_blocks counts block ordinals up to the capacity, including the + # stale ones a sliding window left without a page, so these are logical + # blocks, not pages held. capacity = sum(c.capacity for c in self.kv_cache_map.values()) blocks = sum(int(c.num_blocks) for c in self.kv_cache_map.values()) per_cache = ", ".join( - f"{req_id}: capacity={c.capacity} history={c.history_length} blocks={c.num_blocks}" + f"{req_id}: capacity={c.capacity} history={c.history_length} " + f"logical_blocks={c.num_blocks}" for req_id, c in self.kv_cache_map.items() ) return ( f" [draft pool: {len(self.kv_cache_map)} live caches, capacity {capacity} tokens " - f"in {blocks} allocated blocks ({per_cache}), " + f"in {blocks} logical blocks ({per_cache}), " f"windows={self.max_attention_window_vec}, gpu_max_tokens={self._gpu_max_tokens}]" ) @@ -6131,25 +6136,28 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): + draft_len + self.num_extra_kv_tokens ) - # The chunk start is this cache's history: sliding-window layers - # keep pages only from history + 1 - window (AttnLifeCycle:: - # getStaleRange), so a mirror holds its window plus the chunk - # it is computing rather than the whole prompt so far. - # The executor advances the draft history through - # update_context_resources only under joint reuse, and a reused - # target prefix starts the mirror at position > 0 with no pages - # of its own, so the resize carries the history itself. - history = max(kv_cache.history_length, req.context_current_position) - if req.is_disagg_generation_init_state: - # A disaggregated generation worker receives the prompt KV - # instead of prefilling it, and the standalone drafter - # starts this request on a fresh slot (DFlash/DSpark - # prepare): the mirror never holds drafter K/V for the - # transferred prompt. The whole prompt is therefore - # history, so the window trims the mirror to its tail - # plus the generation headroom instead of one prompt-sized - # chunk that a windowed draft pool cannot hold. - history = max(history, req.context_current_position + req.context_chunk_size) + # None keeps the cache's current history, as every drafter + # that reads the draft request view (MTP, Eagle, two-engine) + # needs: that view's position is not where its layers write. + history = None + if not self._mirror_uses_draft_view: + if req.is_disagg_generation_init_state: + # The transferred prompt KV is the target's only. The + # drafter (DFlash/DSpark: the only separate draft pool a + # disaggregated worker keeps) starts this request on a + # fresh slot at context length 0 and addresses its own + # writes and reads from there, so the mirror is sized in + # those slot-local positions: no page for the prompt. + capacity = draft_len + self.num_extra_kv_tokens + elif kv_cache.capacity >= req.context_current_position: + # The chunk start is history for a mirror that already + # holds [0, chunk start): sliding-window layers keep + # pages only from history + 1 - window (AttnLifeCycle:: + # getStaleRange), so the mirror holds its window plus + # the chunk in flight rather than the whole prompt so + # far. A mirror first seen past position 0 is not one: + # its drafter slot stores that chunk from 0. + history = max(kv_cache.history_length, req.context_current_position) if not kv_cache.resize(capacity, history): raise self._mirror_resize_error(req, "context", capacity) self._after_shrink(req.py_request_id, kv_cache) @@ -7061,11 +7069,11 @@ def release_resources( if not success: release_resources(req, free_draft_resources=True) return None - if is_gen and not materialize_history: - draft_kv_cache.enable_swa_scratch_reuse = False # Same history hint as the target cache above: a sliding-window # draft pool then backs only the dummy's window, not its whole - # length. + # length. The hint exceeds the window only for the + # longest-length warmup dummies, which run before any request + # holds a draft page. success = draft_kv_cache.resize(dummy_capacity, history_length=history_hint) if not success: release_resources(req, free_draft_resources=True) diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py index 64afc689c57a..fe729ddd3286 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py @@ -15,6 +15,7 @@ import array import os +from collections.abc import Callable from dataclasses import dataclass, field, replace from types import SimpleNamespace from unittest.mock import Mock, call, patch @@ -43,6 +44,7 @@ from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.speculative.dflash import DFlashWorker from tensorrt_llm.bindings import DataType, SamplingConfig from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE from tensorrt_llm.bindings.executor import KvCacheRetentionConfig @@ -51,7 +53,10 @@ from tensorrt_llm.conversation_params import ConversationParams from tensorrt_llm.llmapi.llm_args import ( BlockReuseConfig, + DecodingBaseConfig, DFlashDecodingConfig, + DraftTargetDecodingConfig, + DSparkDecodingConfig, Eagle3DecodingConfig, KvCacheConfig, MTPDecodingConfig, @@ -1913,79 +1918,75 @@ def resume_and_restore(request_id, current_cache): @pytest.mark.parametrize( - ("mirror_history", "expected_history"), + ( + "uses_draft_view", + "disagg_gen_init", + "position", + "mirror_capacity", + "mirror_history", + "expected_resize", + ), [ - # A fresh mirror for a request whose target prefix was reused: the - # history jumps to the chunk start so no page below the window is - # allocated for [0, position). - (0, 8192), - # A later chunk of the same request: the history follows the chunk. - (4096, 8192), - # Joint reuse already advanced the history past the chunk start; a - # resize may not decrease it. - (8704, 8704), + # External drafter (target view), aggregated chunked prefill. The + # prompt is 8192 + 2 * 1024 tokens, so the 1024-token chunk at 8192 is + # not the last one and carries no draft tokens. + pytest.param(False, False, 0, 0, 0, (1024, 0), id="first_chunk"), + # A later chunk of a mirror that holds [0, 8192): the chunk start is + # history, so a sliding window drops the pages below it. + pytest.param(False, False, 8192, 8192, 4096, (9216, 8192), id="later_chunk"), + # The history is already ahead of the chunk start; a resize never + # lowers it. + pytest.param(False, False, 8192, 8192, 8704, (9216, 8704), id="history_ahead"), + # A mirror first seen past position 0 (a target prefix this worker did + # not compute): its drafter slot stores the chunk from 0, so nothing + # below the chunk is history and the mirror covers [0, chunk end). + pytest.param(False, False, 8192, 0, 0, (9216, None), id="fresh_past_zero"), + # MTP, Eagle and two-engine drafters keep the draft request view: it + # reads position 0 with the whole prompt as its (last) chunk, and the + # resize keeps the cache's history. + pytest.param(True, False, 8192, 8192, 0, (10243, None), id="draft_view"), + # Disaggregated generation init: the drafter starts a fresh slot at + # context length 0, so the mirror holds no page for the transferred + # prompt (10240 tokens), only the draft headroom. + pytest.param(False, True, 0, 0, 0, (3, None), id="disagg_gen_init-target_view"), + pytest.param(True, True, 0, 0, 0, (10243, None), id="disagg_gen_init-draft_view"), ], ) -def test_unpaired_one_engine_draft_mirror_follows_the_target_chunk( - mirror_history: int, expected_history: int +def test_unpaired_draft_mirror_context_resize( + uses_draft_view: bool, + disagg_gen_init: bool, + position: int, + mirror_capacity: int, + mirror_history: int, + expected_resize: tuple, ) -> None: - """A one-engine external drafter's mirror reads the target's position and - chunk, and its context resize carries the chunk start as history.""" - position, chunk, draft_len = 8192, 1024, 3 + """Which request view an unpaired draft mirror reads, and the capacity and + history its context resize asks for.""" + prompt, chunk, draft_len = 8192 + 2 * 1024, 1024, 3 request = LlmRequest( request_id=41, max_new_tokens=4, - input_tokens=[0] * (position + chunk), - sampling_config=SamplingConfig(1), - is_streaming=False, - draft_tokens=[1] * draft_len, - ) - request.state = LlmRequestState.CONTEXT_INIT - request.context_current_position = position - request.context_chunk_size = chunk - assert request.is_last_context_chunk - - kv_cache = Mock(is_active=True, capacity=mirror_history, history_length=mirror_history) - kv_cache.resize.return_value = True - manager = object.__new__(KVCacheManagerV2) - manager.is_draft = True - manager.enable_joint_kv_cache_reuse = False - manager._mirror_uses_draft_view = False - manager.num_extra_kv_tokens = 0 - manager.kv_cache_map = {request.py_request_id: kv_cache} - manager._resume_and_restore = Mock(return_value=True) - - batch = ScheduledRequests() - batch.context_requests_last_chunk.append(request) - manager._prepare_draft_resources(batch) - - kv_cache.resize.assert_called_once_with(position + chunk + draft_len, expected_history) - assert not request.use_draft_model - - -def test_unpaired_two_engine_draft_mirror_reads_the_draft_request_view() -> None: - """The draft view starts at position 0 with the whole prompt as its chunk - until a separate draft engine advances it, so a two-engine mirror sizes - its cache for the prompt at once.""" - position, chunk, draft_len = 8192, 1024, 3 - request = LlmRequest( - request_id=42, - max_new_tokens=4, - input_tokens=[0] * (position + chunk), + input_tokens=[0] * prompt, sampling_config=SamplingConfig(1), is_streaming=False, draft_tokens=[1] * draft_len, ) - request.state = LlmRequestState.CONTEXT_INIT - request.context_current_position = position - request.context_chunk_size = chunk + if disagg_gen_init: + request.state = LlmRequestState.DISAGG_GENERATION_INIT + assert request.is_last_context_chunk + else: + request.state = LlmRequestState.CONTEXT_INIT + request.context_current_position = position + request.context_chunk_size = chunk + assert not request.is_last_context_chunk - kv_cache = Mock(is_active=True, capacity=0, history_length=0) - kv_cache.resize.return_value = True + kv_cache = Mock(is_active=True, capacity=mirror_capacity, history_length=mirror_history) + views = [] + kv_cache.resize.side_effect = lambda *args: views.append(request.use_draft_model) or True manager = object.__new__(KVCacheManagerV2) manager.is_draft = True manager.enable_joint_kv_cache_reuse = False - manager._mirror_uses_draft_view = True + manager._mirror_uses_draft_view = uses_draft_view manager.num_extra_kv_tokens = 0 manager.kv_cache_map = {request.py_request_id: kv_cache} manager._resume_and_restore = Mock(return_value=True) @@ -1994,8 +1995,8 @@ def test_unpaired_two_engine_draft_mirror_reads_the_draft_request_view() -> None batch.context_requests_last_chunk.append(request) manager._prepare_draft_resources(batch) - kv_cache.resize.assert_called_once_with(position + chunk + draft_len, 0) - assert not request.use_draft_model + kv_cache.resize.assert_called_once_with(*expected_resize) + assert views == [uses_draft_view] @pytest.mark.parametrize("phase", ["context", "generation"]) @@ -2080,7 +2081,10 @@ def test_exhausted_pool_fails_explicitly_on_an_empty_fresh_mirror() -> None: request.context_current_position = 0 request.context_chunk_size = prompt - occupant = Mock(is_active=True, capacity=1216, history_length=0, num_blocks=38) + # Two occupants with distinct fields, so a swapped field or a max in + # place of the sum changes the message. + occupant = Mock(is_active=True, capacity=1216, history_length=512, num_blocks=38) + small = Mock(is_active=True, capacity=64, history_length=0, num_blocks=2) fresh = Mock(is_active=True, capacity=0, history_length=0, num_blocks=0) fresh.resize.return_value = False manager = object.__new__(KVCacheManagerV2) @@ -2088,7 +2092,7 @@ def test_exhausted_pool_fails_explicitly_on_an_empty_fresh_mirror() -> None: manager.enable_joint_kv_cache_reuse = False manager._mirror_uses_draft_view = False manager.num_extra_kv_tokens = 0 - manager.kv_cache_map = {60: occupant, 61: fresh} + manager.kv_cache_map = {60: occupant, 62: small, 61: fresh} manager.max_attention_window_vec = [2048] manager._gpu_max_tokens = 1216 manager._resume_and_restore = Mock(return_value=True) @@ -2104,8 +2108,9 @@ def test_exhausted_pool_fails_explicitly_on_an_empty_fresh_mirror() -> None: f"context resize failed for request 61: could not resize to {prompt + draft_len} tokens" in message ) - assert "2 live caches, capacity 1216 tokens in 38 allocated blocks" in message - assert "61: capacity=0 history=0 blocks=0" in message + assert "3 live caches, capacity 1280 tokens in 40 logical blocks" in message + assert "60: capacity=1216 history=512 logical_blocks=38" in message + assert "61: capacity=0 history=0 logical_blocks=0" in message assert "windows=[2048]" in message assert "gpu_max_tokens=1216" in message @@ -2166,9 +2171,6 @@ def test_windowed_draft_pool_backs_only_the_window_of_a_generation_dummy() -> No dtype=DataType.HALF, spec_config=spec, ) - # Only the external drafter's own manager follows the target's chunking. - assert not draft._mirror_uses_draft_view - assert target._mirror_uses_draft_view stream = torch.cuda.current_stream().cuda_stream def live_pages(cache: object) -> int: @@ -2207,6 +2209,365 @@ def live_pages(cache: object) -> int: target.shutdown() +def _small_spec_manager( + spec: DecodingBaseConfig, + is_draft: bool, + *, + window: int | None = None, + num_kv_heads: int = 2, + dtype: DataType = DataType.HALF, + max_gpu_total_bytes: int = 256 << 20, + max_seq_len: int = 16384, + max_num_tokens: int = 8192, +) -> KVCacheManagerV2: + """A two-layer, head_dim 64, 32-token-page V2 manager for ``spec``, + optionally with a uniform sliding window and no block reuse.""" + num_layers = 2 + config = KvCacheConfig( + enable_block_reuse=False, + max_attention_window=None if window is None else [window] * num_layers, + max_gpu_total_bytes=max_gpu_total_bytes, + ) + return KVCacheManagerV2( + config, + CacheType.SELF, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=64, + tokens_per_block=32, + max_seq_len=max_seq_len, + max_batch_size=2, + max_num_tokens=max_num_tokens, + mapping=Mapping(), + dtype=dtype, + spec_config=spec, + is_draft=is_draft, + ) + + +def _backed_ordinals(cache: object) -> set[int]: + """Block ordinals of ``cache`` that have a page behind them.""" + return { + ordinal + for ordinal, page in enumerate(cache.get_base_page_indices(0)) + if page != BAD_PAGE_INDEX + } + + +@pytest.mark.parametrize( + ("make_spec", "external_drafter"), + [ + pytest.param( + lambda: DFlashDecodingConfig(max_draft_len=4, speculative_model="draft"), + True, + id="dflash", + ), + pytest.param( + lambda: DSparkDecodingConfig(max_draft_len=4, speculative_model="draft"), + True, + id="dspark", + ), + pytest.param( + lambda: PARDDecodingConfig(max_draft_len=4, speculative_model="draft"), + True, + id="pard", + ), + pytest.param( + lambda: DraftTargetDecodingConfig(max_draft_len=4, speculative_model="draft"), + True, + id="draft_target_one_model", + ), + pytest.param(lambda: MTPDecodingConfig(num_nextn_predict_layers=1), False, id="mtp"), + pytest.param( + lambda: Eagle3DecodingConfig(max_draft_len=3, speculative_model="draft"), + False, + id="eagle3_one_model", + ), + ], +) +def test_only_an_external_drafters_draft_manager_reads_the_target_view( + make_spec: Callable[[], DecodingBaseConfig], external_drafter: bool +) -> None: + """The draft manager of a one-engine external drafter reads the target's + position and chunk; every other manager keeps the draft request view.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + spec = make_spec() + assert spec.spec_dec_mode.is_external_drafter() is external_drafter + for is_draft in (False, True): + manager = _small_spec_manager(spec, is_draft, max_gpu_total_bytes=4 << 20) + try: + assert manager._mirror_uses_draft_view is not (is_draft and external_drafter) + finally: + manager.shutdown() + + +@pytest.mark.parametrize( + ("make_spec", "releases_below_window"), + [ + # MTP/Eagle draft layers process the target's chunks through the draft + # pool but the mirror reads the draft request view, which nothing + # advances: it keeps the whole prompt backed. + pytest.param( + lambda: Eagle3DecodingConfig(max_draft_len=3, speculative_model="draft"), + False, + id="eagle3_one_model", + ), + # An external drafter follows the target's chunks and drops what falls + # out of the window. + pytest.param( + lambda: DFlashDecodingConfig(max_draft_len=4, speculative_model="draft"), + True, + id="dflash", + ), + ], +) +def test_windowed_draft_mirror_backs_what_the_next_chunk_attends( + make_spec: Callable[[], DecodingBaseConfig], releases_below_window: bool +) -> None: + """Chunked prefill with block reuse off on a sliding-window draft pool: + after the second chunk's resize, every page the chunk attends (its window + back from the chunk start, through the chunk end) is backed.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + window, page, prompt, chunk = 256, 32, 8192, 512 + draft = _small_spec_manager(make_spec(), True, window=window) + request = LlmRequest( + request_id=101, + max_new_tokens=8, + input_tokens=[1] * prompt, + sampling_config=SamplingConfig(1), + is_streaming=False, + ) + request.state = LlmRequestState.CONTEXT_INIT + try: + request.context_chunk_size = chunk + batch = ScheduledRequests() + batch.context_requests_chunking.append(request) + draft._prepare_draft_resources(batch) + + request.move_to_next_context_chunk() + request.context_chunk_size = chunk + assert request.context_current_position == chunk + batch = ScheduledRequests() + batch.context_requests_chunking.append(request) + draft._prepare_draft_resources(batch) + + backed = _backed_ordinals(draft.kv_cache_map[request.py_request_id]) + attended = set(range((chunk - window) // page, 2 * chunk // page)) + assert attended <= backed + assert (0 in backed) is not releases_below_window + draft.free_resources(request) + finally: + draft.shutdown() + + +class _PoolBoundDFlashDraftModel: + """Just enough of a DFlash draft model for the worker to bind the draft + KV pool and run its generation-step context store.""" + + block_size = 4 + config = SimpleNamespace(max_position_embeddings=1 << 17) + dflash_attention_backend = "TRTLLM" + hidden_norm = object() + + def __init__(self, num_layers: int, num_kv_heads: int, head_dim: int, hidden: int) -> None: + self.fc = SimpleNamespace(weight=torch.empty(1, dtype=torch.bfloat16, device="cuda")) + self._layers = (num_layers, num_kv_heads, head_dim) + embed_tokens = torch.nn.Embedding(16, hidden).cuda() + self.draft_model_full = SimpleNamespace(model=SimpleNamespace(embed_tokens=embed_tokens)) + + def _build_fused_kv_buffers(self) -> None: + self._num_attn_layers, self._num_kv_heads, self._head_dim = self._layers + self._num_heads = self._num_kv_heads + + def _get_attention_mask_args(self, layer_idx: int) -> tuple: + return True, (-1, -1) + + def validate_block_attention_windows(self) -> None: + pass + + def project_target_hidden(self, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states + + def precompute_context_kv( + self, projected: torch.Tensor, positions: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + num_layers, num_kv_heads, head_dim = self._layers + k = torch.full( + (projected.shape[0], num_layers, num_kv_heads, head_dim), + 3.0, + dtype=torch.bfloat16, + device="cuda", + ) + return k, k.clone() + + +def test_disagg_generation_drafter_writes_stay_on_its_own_draft_pages( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A disaggregated generation worker receives a prompt far larger than its + sliding-window draft pool. The mirror must fit, and the drafter's first + generation step -- a fresh slot that stores the accepted tokens' context + K/V from position 0, then appends its block -- must write only pages its + own mirror holds, never one another request owns.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + monkeypatch.setattr( + "tensorrt_llm._torch.speculative.dflash.validate_dflash_trtllm_gen_runtime", + lambda **kwargs: None, + ) + num_layers, num_kv_heads, head_dim, hidden = 2, 1, 64, 8 + window, prompt, occupant_prompt, max_draft_len = 256, 65536, 512, 3 + bytes_per_token = num_layers * 2 * num_kv_heads * head_dim * 2 + spec = DFlashDecodingConfig( + max_draft_len=max_draft_len, + speculative_model="draft", + attention_backend="TRTLLM", + mask_token_id=0, + ) + # Room for 8192 tokens: the occupant and a windowed mirror fit, a mirror + # of the whole transferred prompt does not. + draft = _small_spec_manager( + spec, + True, + window=window, + num_kv_heads=num_kv_heads, + dtype=DataType.BF16, + max_gpu_total_bytes=8192 * bytes_per_token, + max_seq_len=1 << 17, + max_num_tokens=occupant_prompt, + ) + occupant = LlmRequest( + request_id=1, + max_new_tokens=8, + input_tokens=[1] * occupant_prompt, + sampling_config=SamplingConfig(1), + is_streaming=False, + ) + occupant.state = LlmRequestState.CONTEXT_INIT + occupant.context_chunk_size = occupant_prompt + request = LlmRequest( + request_id=2, + max_new_tokens=8, + input_tokens=[1] * prompt, + sampling_config=SamplingConfig(1), + is_streaming=False, + draft_tokens=[1] * max_draft_len, + ) + request.state = LlmRequestState.DISAGG_GENERATION_INIT + offsets = torch.zeros( + (getattr(draft, "num_attention_op_pools", draft.num_pools), 2, 2, draft.max_blocks_per_seq), + dtype=torch.int32, + device="cuda", + ) + try: + batch = ScheduledRequests() + batch.context_requests_last_chunk.append(occupant) + draft._prepare_draft_resources(batch) + batch = ScheduledRequests() + batch.context_requests_last_chunk.append(request) + draft._prepare_draft_resources(batch) + request.state = LlmRequestState.GENERATION_IN_PROGRESS + batch = ScheduledRequests() + batch.generation_requests.append(request) + draft._prepare_draft_resources(batch) + + worker = DFlashWorker(spec, Mapping()) + draft_model = _PoolBoundDFlashDraftModel(num_layers, num_kv_heads, head_dim, hidden) + worker.set_draft_model(draft_model) + attn_metadata = SimpleNamespace( + max_seq_len=1 << 17, + draft_kv_cache_block_offsets=offsets, + num_contexts=0, + num_seqs=1, + num_ctx_tokens=0, + _seq_lens_cuda=torch.zeros(1, dtype=torch.int32, device="cuda"), + _seq_lens=torch.zeros(1, dtype=torch.int32), + ) + worker._lazy_init_ctx_buffers( + draft_model, + SimpleNamespace(max_num_requests=2), + attn_metadata, + draft_kv_cache_manager=draft, + ) + assert worker._ctx_block_tables is not None, "the drafter did not bind the draft pool" + # V2 hands out per-layer views over one interleaved pool: a table entry + # names layer 0's slot of a page, and layer l of that page is slot + # entry + l of the same dense view, so view 0 sees every layer's write. + assert worker._ctx_block_divisor == draft.kv_factor + pool = worker._ctx_kv_buf[0] + + def own_slots(req: LlmRequest) -> set[int]: + draft.copy_batch_block_offsets(offsets, [req.py_request_id], 1, 0, 1) + worker._refresh_ctx_block_tables(attn_metadata, 1) + row = worker._ctx_block_tables[0].tolist() + return { + row[ordinal] + layer + for ordinal in _backed_ordinals(draft.kv_cache_map[req.py_request_id]) + for layer in range(num_layers) + } + + occupant_slots = own_slots(occupant) + request_slots = own_slots(request) + assert occupant_slots and request_slots + assert not occupant_slots & request_slots + + # As DFlashSpecMetadata.prepare() does for a new generation request. + slot = worker._assign_slot(request.py_request_id) + worker._rebuild_batch_to_slot([request.py_request_id]) + assert worker._ctx_len_host[slot] == 0 + pool.fill_(7.0) + before = pool.clone() + + tokens_per_req = max_draft_len + 1 + captured = torch.randn(tokens_per_req, hidden, device="cuda") + spec_metadata = SimpleNamespace( + hidden_size=hidden, + runtime_draft_len=max_draft_len, + captured_hidden_states=None, + get_hidden_states=lambda num_tokens: captured[:num_tokens], + ) + inputs = worker.prepare_1st_drafter_inputs( + input_ids=torch.zeros(tokens_per_req, dtype=torch.long, device="cuda"), + position_ids=torch.zeros(tokens_per_req, dtype=torch.long, device="cuda"), + hidden_states=torch.zeros(tokens_per_req, hidden, device="cuda"), + accepted_tokens=torch.ones((1, tokens_per_req), dtype=torch.long, device="cuda"), + num_accepted_tokens=torch.tensor([tokens_per_req], dtype=torch.int32, device="cuda"), + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + total_target_tokens=tokens_per_req, + ) + # The block decode then appends the block's own K/V right after the + # context it reads, through the same table row. + num_ctx = int(inputs["num_ctx_per_req"][0]) + assert num_ctx == tokens_per_req + block = worker._compute_block_size + block_kv = torch.full( + (block, num_layers, num_kv_heads, head_dim), 5.0, dtype=torch.bfloat16, device="cuda" + ) + worker._store_context_kv_paged( + block_kv, + block_kv.clone(), + torch.zeros(block, dtype=torch.long, device="cuda"), + torch.arange(num_ctx, num_ctx + block, dtype=torch.long, device="cuda"), + ) + torch.cuda.synchronize() + + written = set((pool != before).flatten(1).any(dim=1).nonzero().flatten().tolist()) + assert written, "the drafter wrote nothing" + assert not written & occupant_slots, "the drafter wrote another request's draft pages" + assert written <= request_slots + draft.free_resources(request) + draft.free_resources(occupant) + finally: + draft.shutdown() + + def _make_publishing_manager(policy: BlockReusePolicy) -> KVCacheManagerV2: """A manager wired just far enough to run the publish/history bookkeeping.""" manager = object.__new__(KVCacheManagerV2) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_kv_resume_pressure.py b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_kv_resume_pressure.py index 2daeb957320d..01e1edd11e67 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_kv_resume_pressure.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_kv_resume_pressure.py @@ -45,7 +45,7 @@ class _Req: - def __init__(self, req_id): + def __init__(self, req_id: int) -> None: self.py_request_id = req_id self.lora_task_id = None self.cache_salt = None @@ -68,7 +68,7 @@ def all_requests(self): return self.context_requests + self.generation_requests -def _manager(resume_ok): +def _manager(resume_ok: bool) -> tuple[KVCacheManagerV2, MagicMock]: """A draft manager whose resume answers ``resume_ok``, with nothing else real. Built without ``__init__``: the method under test needs only the mirror @@ -128,39 +128,6 @@ def test_context_request_is_prepared_normally_when_resume_succeeds(): cache.resize.assert_called_once() -def test_disagg_generation_init_sizes_the_mirror_from_history_and_chunk() -> None: - """A transferred prompt is history for the drafter's mirror, not a chunk to hold. - - On a disaggregated generation worker the prompt KV arrives over the wire - and the one-engine drafter starts the request on a fresh slot, so the - mirror never holds drafter K/V for that prompt. The resize therefore - carries the whole prompt as history, which lets the attention window trim - the mirror to its tail plus the generation headroom; a history of 0 asks - a windowed draft pool for one prompt-sized chunk instead. - """ - manager, cache = _manager(resume_ok=True) - req = _Req(9) - req.is_disagg_generation_init_state = True - req.context_current_position = 0 - req.context_chunk_size = 200_000 - - manager._prepare_draft_resources(_Batch(context_requests=[req])) - - cache.resize.assert_called_once_with(200_000, 200_000) - - -def test_ordinary_context_chunk_keeps_its_start_as_history() -> None: - """The aggregated chunked path: history is the chunk start, not its end.""" - manager, cache = _manager(resume_ok=True) - req = _Req(10) - req.context_current_position = 64 - req.context_chunk_size = 8 - - manager._prepare_draft_resources(_Batch(context_requests=[req])) - - cache.resize.assert_called_once_with(72, 64) - - def test_generation_request_still_raises_on_a_resume_refusal(): """Skipping a generation request asserts later in C++, so it must raise here.""" manager, _ = _manager(resume_ok=False) From f8cbd7f164a8155fc49179f897eef532556bfea0 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Thu, 29 Oct 2026 17:58:44 -0700 Subject: [PATCH 03/14] [TRTLLM-17390][fix] Keep the drafter and its draft mirror in absolute positions A disaggregated generation-init mirror is sized in absolute positions again, with the transferred prompt as history. DFlash seeds a fresh generation slot on the page holding its absolute position and reads through a per-slot rebased block table, so a slot never addresses pages below what it stored. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- .../spec_worker.py | 3 + .../spec_worker.py | 3 + .../kv_cache/kv_cache_manager_v2.py | 25 +- tensorrt_llm/_torch/speculative/dflash.py | 319 +++++++++++++++--- .../_torch/speculative/spec_tail_kernels.py | 36 +- .../kv_cache/test_kvcm2_integration.py | 227 +++++++++---- .../hw_agnostic/test_dflash_copy_skip.py | 4 +- .../hw_agnostic/test_dflash_draft_kv.py | 280 +++++++++++++++ .../hw_agnostic/test_dflash_dummy_slot.py | 6 +- .../test_dflash_guided_rollback.py | 2 + .../speculative/test_spec_tail_kernels.py | 34 +- 11 files changed, 807 insertions(+), 132 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/spec_worker.py index 277d1f07ecf3..a35aa480082e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/spec_worker.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/spec_worker.py @@ -61,6 +61,9 @@ class KimiK3DFlashWorker(DFlashWorker): # inputs as trtllm::k3_spec_accept, keep the target logits vocabulary-sharded and write the drafter's context # K/V with trtllm::k3_ctx_kv. k3_decode = False + # The generation-step preparation below reads _ctx_block_tables, not the + # rebased view, so every drafter slot numbers from 0. + _ctx_reads_rebased_view = False def _mask_token_id(self, draft_model) -> int: """The drafter's mask token id, resolved once.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py index 277d1f07ecf3..a35aa480082e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py @@ -61,6 +61,9 @@ class KimiK3DFlashWorker(DFlashWorker): # inputs as trtllm::k3_spec_accept, keep the target logits vocabulary-sharded and write the drafter's context # K/V with trtllm::k3_ctx_kv. k3_decode = False + # The generation-step preparation below reads _ctx_block_tables, not the + # rebased view, so every drafter slot numbers from 0. + _ctx_reads_rebased_view = False def _mask_token_id(self, draft_model) -> int: """The drafter's mask token id, resolved once.""" diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 4cf63a687a8d..dd234f3bf8ce 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -6142,21 +6142,28 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): history = None if not self._mirror_uses_draft_view: if req.is_disagg_generation_init_state: - # The transferred prompt KV is the target's only. The - # drafter (DFlash/DSpark: the only separate draft pool a - # disaggregated worker keeps) starts this request on a - # fresh slot at context length 0 and addresses its own - # writes and reads from there, so the mirror is sized in - # those slot-local positions: no page for the prompt. - capacity = draft_len + self.num_extra_kv_tokens + # The transferred prompt KV is the target's only, so + # the mirror holds no drafter K/V below the prompt end. + # The mirror and the drafter share absolute positions: + # the drafter (DFlash/DSpark, the only separate draft + # pool a disaggregated worker keeps) seeds this + # request's slot at the prompt end and reads nothing + # below it (DFlashSpecMetadata.prepare). The whole + # prompt is therefore history, which a sliding window + # trims to its tail plus the generation headroom, and + # which update_resources keeps advancing absolutely. + history = max( + kv_cache.history_length, + req.context_current_position + req.context_chunk_size, + ) elif kv_cache.capacity >= req.context_current_position: # The chunk start is history for a mirror that already # holds [0, chunk start): sliding-window layers keep # pages only from history + 1 - window (AttnLifeCycle:: # getStaleRange), so the mirror holds its window plus # the chunk in flight rather than the whole prompt so - # far. A mirror first seen past position 0 is not one: - # its drafter slot stores that chunk from 0. + # far. A mirror first seen past position 0 is not one + # and keeps its history, i.e. pages from 0. history = max(kv_cache.history_length, req.context_current_position) if not kv_cache.resize(capacity, history): raise self._mirror_resize_error(req, "context", capacity) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 2eb3965327ee..733da12b36a1 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -17,7 +17,7 @@ import weakref from collections import deque from dataclasses import dataclass -from typing import TYPE_CHECKING, Callable, List, Optional +from typing import TYPE_CHECKING, Callable, List, Optional, Tuple import torch from torch import nn @@ -516,6 +516,10 @@ def dflash_allocated_ctx_limit( ) -> torch.Tensor: """Context length a request may advertise, given the pages it holds. + ``block_counts`` is each row's extent: one past its last backed entry + (``decode_draft_block_table``), so positions are absolute even when a + sliding window released the low entries. + Leaves ``block_size`` positions of room: the MLA path writes the block's own latents at ``ctx_len .. ctx_len + block_size``, so stopping at exactly the allocation puts the first of them on an unallocated page, whose table entry @@ -531,6 +535,51 @@ def dflash_allocated_ctx_limit( return (block_counts * page_size - block_size).clamp(min=0) +def dflash_windows_span_a_page(windows: Optional[List[Optional[int]]], page_size: int) -> bool: + """Whether every attention window of a draft pool spans at least a page. + + ``None`` (for the list or an entry) is a layer without a window, which + keeps every page it was given. + """ + return all(window >= page_size for window in windows or () if window is not None) + + +def dflash_rebased_ctx_rows(rows: torch.Tensor, base_pages: torch.Tensor) -> torch.Tensor: + """Block-table rows shifted left by ``base_pages`` entries each. + + Entry ``j`` of the result names the page entry ``j + base`` of the input + row named, so a reader that numbers positions from ``base * page_size`` + reaches the same physical pages as one numbering from 0 and never + addresses an entry below ``base``. Entries shifted in from past the row's + end repeat the last entry; a reader never dereferences them because its + sequence length is shortened by the same ``base * page_size``. + """ + width = rows.size(1) + index = torch.arange(width, device=rows.device).unsqueeze(0) + base_pages.view(-1, 1) + return rows.gather(1, index.clamp_(max=width - 1)) + + +def dflash_rebased_ctx_view( + rows: torch.Tensor, base_pages: torch.Tensor, num_ctx: torch.Tensor, page_size: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """Block-table rows and context lengths the block decode reads for slots + whose first ``base_pages`` pages hold no drafter K/V (see + ``DFlashWorker._ctx_base_pages``). + + The context length is shortened by the base, so the decode never reads + below it. The decode also writes the block's own K/V through the returned + rows, at rebased positions ``num_ctx .. num_ctx + block_size``, so entry 0 + of a row has to name a page the request holds. While the context is + shorter than the base (the first tokens stored inside the rounded-up page) + it reads nothing and the row is shifted only to the page holding the + context length, which is the page those tokens were written to. From + ``base * page_size`` on the shift is the base itself. + """ + rebase_pages = torch.minimum(base_pages, num_ctx // page_size) + num_ctx = (num_ctx - base_pages * page_size).clamp(min=0) + return dflash_rebased_ctx_rows(rows, rebase_pages), num_ctx + + @dataclass class DFlashSpecMetadata(SpecMetadata): """Metadata for DFlash speculative decoding. @@ -602,32 +651,48 @@ def prepare(self): worker = getattr(self, "_dflash_worker", None) if worker is not None and worker._ctx_buf_inited: current = set(self.request_ids) - evicted = {} + # Slot lengths (and bases) written in one batch: evicted slots are + # emptied, newly assigned ones seeded below. + updates, bases = {}, {} for rid in list(worker._req_to_slot.keys()): if rid not in current: slot = worker._req_to_slot.pop(rid) - evicted[slot] = 0 + updates[slot] = 0 worker._req_ctx_pos.pop(rid, None) worker._free_slots.append(slot) - worker._write_ctx_len(evicted) # A disagg generation worker receives prompt KV instead of # prefilling, so _store_prefill_context -- the only other assigner -- - # never runs and every request would land on the dummy slot. + # never runs and every request would land on the dummy slot. The + # same holds for a request whose slot was evicted above. Each such + # slot starts empty at its row's absolute position, on a page + # boundary (_seed_generation_slot). num_contexts = max(0, len(self.request_ids) - self.num_generations) - for rid in self.request_ids[num_contexts:]: + row_starts = None + for row, rid in enumerate(self.request_ids[num_contexts:], start=num_contexts): if ( rid != ATTENTION_DP_DUMMY_REQUEST_ID and rid < worker._graph_dummy_id_floor and rid not in worker._req_to_slot ): - if worker._assign_slot(rid) is None: + slot = worker._assign_slot(rid, updates=updates, bases=bases) + if slot is None: # Not fatal -- spec dec is lossless, so the dummy-slot # fallback only costs acceptance. Warn so it is not silent. logger.warning( f"DFlash: no free drafter slot for generation request {rid}; " "it shares the dummy slot and will draft poorly." ) + continue + if self.host_position_ids is None or self.seq_lens is None: + continue + if row_starts is None: + row_starts = [0] + for seq_len in self.seq_lens: + row_starts.append(row_starts[-1] + int(seq_len)) + position = int(self.host_position_ids[row_starts[row]]) + worker._seed_generation_slot(slot, position, updates, bases) + worker._write_ctx_len(updates, bases) # Route unknown request IDs (cuda graph padding or warmup dummies) # to dummy slot to avoid corrupting real request's context @@ -735,6 +800,14 @@ class DFlashWorker(SpecWorkerBase): Reference: https://arxiv.org/pdf/2602.06036 """ + # Whether the block decode reads _ctx_block_tables_read, the per-slot view + # rebased past _ctx_base_pages. Only then can a slot on the managed pool + # start at an absolute position other than 0 (_ctx_store_start, + # _seed_generation_slot); a subclass with its own generation-step + # preparation that reads _ctx_block_tables directly sets it False and + # numbers every slot from 0. + _ctx_reads_rebased_view = True + def __init__( self, spec_config: "DFlashDecodingConfig", @@ -789,9 +862,22 @@ def __init__( # manager it is replaced by _ctx_block_tables, refreshed per iteration. self._ctx_page_table = None self._ctx_block_tables = None # [max_batch+1, max_blocks_per_seq] int32 + # The block table the block decode reads: each row of _ctx_block_tables + # shifted left by that slot's _ctx_base_pages, so pages below the first + # position the drafter stored are not addressable. + self._ctx_block_tables_read = None + # [max_batch+1] pages below a slot's first stored position: the slot + # holds no drafter K/V there (a prompt this worker did not prefill, or + # a prefix the target reused from an unpaired pool). Host shadow and + # pre-forward snapshots as for _ctx_len. + self._ctx_base_pages = None + self._ctx_base_host = None + self._saved_ctx_base = None + self._saved_ctx_base_host = None self._ctx_block_indptr = None - # Pages the manager handed this request, per batch row. Saturates at - # the table width on the default backend -- dflash_allocated_ctx_limit. + # Extent of the pages the manager handed this request (one past its + # last backed entry), per batch row. Saturates at the table width on + # the default backend -- dflash_allocated_ctx_limit. self._ctx_block_counts = None self._ctx_pool_idx = 0 self._ctx_block_divisor = 1 # encoded offset -> raw pool block index @@ -805,6 +891,11 @@ def __init__( # makes a reused prefix's drafter K/V still addressable. Set in # _lazy_init_ctx_buffers; read by _store_prefill_context. self._ctx_reuse_addressable = False + # Whether a seeded slot may start at the first position of the page + # holding its row's position (_seed_generation_slot): true when every + # window of the bound draft pool spans at least a page. Set in + # _lazy_init_ctx_buffers. + self._ctx_seed_page_aligned = False self._dflash_attention_backend = spec_config.attention_backend # Set by bind_target_lm_head. self._vocab_parallel: Optional[VocabParallelSelection] = None @@ -1146,6 +1237,7 @@ def _init_ctx_block_tables( self._ctx_block_tables = torch.zeros( (num_slots, max_blocks), dtype=torch.int32, device="cuda" ) + self._ctx_block_tables_read = torch.zeros_like(self._ctx_block_tables) self._ctx_block_indptr = torch.arange( 0, (num_slots + 1) * max_blocks, max_blocks, dtype=torch.int32, device="cuda" ) @@ -1216,6 +1308,7 @@ def _lazy_init_ctx_buffers( # pass falling back to the arena still rebinds once the real pool lands. self._ctx_kv_manager = draft_kv_cache_manager self._ctx_block_tables = None + self._ctx_block_tables_read = None self._ctx_block_counts = None self._paged_latent_prepare = False max_batch = spec_metadata.max_num_requests @@ -1272,6 +1365,8 @@ def _lazy_init_ctx_buffers( self._ctx_len = torch.zeros(num_slots, dtype=torch.long, device="cuda") self._ctx_len_host = [0] * num_slots + self._ctx_base_pages = torch.zeros(num_slots, dtype=torch.long, device="cuda") + self._ctx_base_host = [0] * num_slots self._batch_to_slot_buffer = DeviceCopyBuffer( torch.zeros(max_batch, dtype=torch.long, device="cuda"), pin_memory=prefer_pinned(), @@ -1402,6 +1497,9 @@ def _lazy_init_ctx_buffers( self._ctx_reuse_addressable = bool( getattr(draft_kv_cache_manager, "enable_joint_kv_cache_reuse", False) ) + self._ctx_seed_page_aligned = dflash_windows_span_a_page( + getattr(draft_kv_cache_manager, "max_attention_window_vec", None), page_size + ) self._ctx_kv_last_page_len = torch.full( (num_slots,), page_size, dtype=torch.int32, device="cuda" ) @@ -1454,13 +1552,93 @@ def _ctx_paged_index_args(self): indptr = self._ctx_kv_indptr return table.flatten().contiguous(), indptr - def _assign_slot(self, req_id: int, reset: bool = False, updates: Optional[dict] = None): + def _ctx_store_start(self, cur: int, first_pos: int) -> int: + """Position at which a context chunk's drafter K/V is stored. + + ``cur`` is the slot's stored length, ``first_pos`` the chunk's first + target position. They differ only when the target skipped + [0, first_pos) (block reuse, or a prefix computed elsewhere), so just + the tail reaches the capture buffer. + + On the managed pool the block table is indexed by target position, so + the tail is stored at ``first_pos``. The private arena numbers every + request from 0, where the tail is all the context there is. + """ + if ( + cur == 0 + and first_pos > 0 + and self._ctx_block_tables is not None + and self._ctx_reads_rebased_view + ): + return first_pos + return cur + + def _ctx_base_for_store(self, cur: int, start: int) -> Optional[int]: + """Pages the block decode must not read for a slot whose first stored + position is ``start``, or None when the slot numbers from ``cur``. + + A pool paired with the target's reuse holds the prefix's drafter K/V + below ``start``, so it stays readable. An unpaired pool holds no drafter + K/V there: the block decode is handed the slot's block-table row shifted + past those pages (``dflash_rebased_ctx_rows``) and a context length + shortened by the same amount, so a tail shorter than the drafter's + window attends only to itself. ``start`` is rounded up to a page + boundary; the tail tokens inside that page stay stored but unread. + """ + if start == cur or self._ctx_reuse_addressable: + return None + return -(-start // self._ctx_page_size) + + def _seed_generation_slot(self, slot: int, position: int, updates: dict, bases: dict) -> None: + """Record the context length and base of a slot newly assigned to a + generation row whose first token sits at absolute ``position``. + + Bound to the draft KV cache manager, the slot addresses the request's + mirror, whose positions are absolute (the target's). The slot holds no + drafter K/V for the positions before ``position`` -- a disaggregated + generation worker received only the target's prompt KV, and a request + that lost its slot lost its context -- whether or not the pool joins + block reuse: the mirror of such a request never looked a prefix up. + The slot therefore starts with an empty context on a page boundary, + and the block decode reads from its base, never below it: + + * at the first position of the page holding ``position``. The slot's + positions then trail the absolute ones by less than a page, so they + stay on pages the mirror backs as long as every window of the pool + spans a page (``_ctx_seed_page_aligned``), and the slot reads what + it stores from its first step; + * otherwise at ``position`` itself, with the base rounded up to the + next page: the slot reads nothing until its context reaches it. + + The private arena is slot-addressed and starts at 0, as does a worker + that does not read the rebased view (_ctx_reads_rebased_view). + """ + if self._ctx_block_tables is None or not self._ctx_reads_rebased_view or position <= 0: + updates[slot] = 0 + return + position = min(position, self._max_ctx) + page = self._ctx_page_size + if self._ctx_seed_page_aligned: + bases[slot] = position // page + updates[slot] = bases[slot] * page + else: + bases[slot] = -(-position // page) + updates[slot] = position + + def _assign_slot( + self, + req_id: int, + reset: bool = False, + updates: Optional[dict] = None, + bases: Optional[dict] = None, + ) -> Optional[int]: """Get (or refresh) this request's context slot; None when none is free. ``reset`` recycles an existing slot, which is how a reused request id is - told apart from a continuation. Pass ``updates`` to defer the length - resets into a caller's batch instead of writing _ctx_len per request; - either way the host mirror stays in step with the device tensor. + told apart from a continuation. Pass ``updates`` (and ``bases``, which + then gets the base-page resets) to defer the length resets into a + caller's batch instead of writing _ctx_len per request; either way the + host mirror stays in step with the device tensor. """ def clear(slot: int) -> None: @@ -1468,6 +1646,8 @@ def clear(slot: int) -> None: self._write_ctx_len({slot: 0}) else: updates[slot] = 0 + if bases is not None: + bases[slot] = 0 if reset and req_id in self._req_to_slot: old_slot = self._req_to_slot.pop(req_id) @@ -1590,30 +1770,51 @@ def _rebuild_batch_to_slot(self, request_ids: List[int]) -> None: mapping = [self._req_to_slot.get(rid, self._dummy_slot) for rid in request_ids] self._batch_to_slot_buffer.update(torch.tensor(mapping, dtype=torch.long)) - def _write_ctx_len(self, updates: dict[int, int]) -> None: - """Apply a slot-to-length mapping to _ctx_len in one async scatter.""" - if not updates: + def _write_ctx_len( + self, updates: dict[int, int], base_updates: Optional[dict[int, int]] = None + ) -> None: + """Apply slot-to-length (and slot-to-base-page) mappings in one async + scatter each. A slot reset to length 0 also drops its base pages, and + that reset wins over a base recorded for the same slot in + ``base_updates`` (a chunk's base is recorded before the overflow check + that may reset the slot).""" + if not updates and not base_updates: return + bases = dict(base_updates or {}) + bases.update({slot: 0 for slot, value in updates.items() if value == 0}) for slot, value in updates.items(): # Update host side mirror self._ctx_len_host[slot] = value + for slot, value in bases.items(): + self._ctx_base_host[slot] = value - slots, values = zip(*updates.items()) pinned = prefer_pinned() - slots_pinned = torch.tensor(slots, dtype=torch.long, pin_memory=pinned) - vals_pinned = torch.tensor(values, dtype=torch.long, pin_memory=pinned) - self._ctx_len.index_copy_( - 0, - slots_pinned.to("cuda", non_blocking=True), - vals_pinned.to("cuda", non_blocking=True), - ) + for target, mapping in ((self._ctx_len, updates), (self._ctx_base_pages, bases)): + if not mapping: + continue + slots, values = zip(*mapping.items()) + slots_pinned = torch.tensor(slots, dtype=torch.long, pin_memory=pinned) + vals_pinned = torch.tensor(values, dtype=torch.long, pin_memory=pinned) + target.index_copy_( + 0, + slots_pinned.to(target.device, non_blocking=True), + vals_pinned.to(target.device, non_blocking=True), + ) def _restore_ctx_len_host(self) -> None: """Roll the host shadows back to their pre-forward snapshot.""" if self._saved_ctx_len_host is not None: self._ctx_len_host = list(self._saved_ctx_len_host) + if self._saved_ctx_base_host is not None: + self._ctx_base_host = list(self._saved_ctx_base_host) if self._saved_req_ctx_pos is not None: self._req_ctx_pos = dict(self._saved_req_ctx_pos) + def _restore_ctx_len_device(self) -> None: + """Roll _ctx_len and _ctx_base_pages back to their pre-forward snapshot.""" + self._ctx_len.copy_(self._saved_ctx_len) + if self._saved_ctx_base is not None: + self._ctx_base_pages.copy_(self._saved_ctx_base) + def _store_prefill_context( self, draft_model, @@ -1653,6 +1854,7 @@ def _store_prefill_context( # the same values, but reading it costs a D2H sync per request. host_pos = spec_metadata.host_position_ids ctx_len_updates = {} + ctx_base_updates = {} offset = 0 num_contexts = attn_metadata.num_contexts # A context write addresses pages through row i of the block table, so @@ -1682,7 +1884,12 @@ def _store_prefill_context( # request id reused after completion, a prefill restarted after # preemption -- has to start from a clean slot. reset = self._req_ctx_pos.get(req_id) != first_pos - if self._assign_slot(req_id, reset=reset, updates=ctx_len_updates) is None: + if ( + self._assign_slot( + req_id, reset=reset, updates=ctx_len_updates, bases=ctx_base_updates + ) + is None + ): logger.warning("DFlash: no free slots, skipping context store") self._req_ctx_pos.pop(req_id, None) offset += slen @@ -1691,11 +1898,14 @@ def _store_prefill_context( slot = self._req_to_slot[req_id] cur = ctx_len_updates.get(slot, self._ctx_len_host[slot]) cap = self._max_ctx if ctx_alloc is None else min(self._max_ctx, ctx_alloc[i]) - if cur == 0 and first_pos > 0 and self._ctx_reuse_addressable: - # Reuse let the target skip [0, first_pos), so only the tail - # reaches the capture buffer. The prefix K/V are still in the - # paired pool: precompute_context_kv is per-token in (hidden, p). - cur = first_pos + # A paired pool still holds the prefix K/V below the tail + # (precompute_context_kv is per-token in (hidden, p)); an unpaired + # one does not, so the slot reads from a base above it. + start = self._ctx_store_start(cur, first_pos) + base = self._ctx_base_for_store(cur, start) + if base is not None: + ctx_base_updates[slot] = base + cur = start if cur + slen > cap: # Request-level, like the no-free-slots path above: truncating # would silently draft from a stale prefix, but killing the @@ -1754,7 +1964,7 @@ def _store_prefill_context( self._ctx_v_buf[slot, :, cur:end] = chunk_v.permute(1, 0, 2, 3) offset += slen - self._write_ctx_len(ctx_len_updates) + self._write_ctx_len(ctx_len_updates, ctx_base_updates) def _ensure_spec_dec_state_restored(self, attn_metadata, spec_metadata): # Restore first (in warmup mode kv_lens_cuda was saved and comes back @@ -1774,7 +1984,7 @@ def _ensure_spec_dec_state_restored(self, attn_metadata, spec_metadata): ): # A failed forward must not keep this iteration's in-place # _ctx_len updates: roll back to the pre-forward snapshot. - self._ctx_len.copy_(self._saved_ctx_len) + self._restore_ctx_len_device() self._restore_ctx_len_host() self._ctx_len_restore_pending = False @@ -1840,6 +2050,8 @@ def _forward_impl( # mutate _ctx_len until replay. self._saved_ctx_len = self._ctx_len.clone() self._saved_ctx_len_host = list(self._ctx_len_host) + self._saved_ctx_base = self._ctx_base_pages.clone() + self._saved_ctx_base_host = list(self._ctx_base_host) self._saved_req_ctx_pos = dict(self._req_ctx_pos) self._ctx_len_restore_pending = True @@ -2043,7 +2255,7 @@ def _forward_impl( # Restore context lengths after warmup; real runs keep the updates. if is_warmup: - self._ctx_len.copy_(self._saved_ctx_len) + self._restore_ctx_len_device() self._restore_ctx_len_host() self._ctx_len_restore_pending = False @@ -2392,6 +2604,7 @@ def prepare_1st_drafter_inputs( self._advance_ctx_len(slots, gen_num_accepted_long) num_ctx_per_req_t = self._ctx_len[slots] + ctx_page_table = self._ctx_page_table if self._ctx_block_tables is not None: # The write above clamps columns to the request's allocation, so # a context that outruns it has its tail written over the last @@ -2401,6 +2614,18 @@ def prepare_1st_drafter_inputs( self._ctx_block_counts, self._ctx_page_size, block_size ) num_ctx_per_req_t = torch.minimum(num_ctx_per_req_t, allocated[gen_rows_out]) + # Pages below a slot's base hold no drafter K/V: the read + # numbers from the base instead (dflash_rebased_ctx_view). + # Writes above keep the absolute table. + self._ctx_block_tables_read[gen_rows_out], num_ctx_per_req_t = ( + dflash_rebased_ctx_view( + self._ctx_block_tables[gen_rows_out], + self._ctx_base_pages[slots], + num_ctx_per_req_t, + self._ctx_page_size, + ) + ) + ctx_page_table = self._ctx_block_tables_read noise_embedding = noise_embed_2d query_positions = query_position_ids.long() @@ -2417,6 +2642,11 @@ def prepare_1st_drafter_inputs( slots = torch.empty(0, dtype=torch.long, device="cuda") gen_rows_out = slots bonus = torch.empty(0, dtype=torch.long, device="cuda") + ctx_page_table = ( + self._ctx_block_tables_read + if self._ctx_block_tables is not None + else self._ctx_page_table + ) return { "lilicorr_anchor": lilicorr_anchor, @@ -2432,11 +2662,8 @@ def prepare_1st_drafter_inputs( # token. The dspark Markov chain conditions its first step on it. "first_prev_tokens": bonus, "ctx_kv_cache": self._ctx_kv_buf, - "ctx_page_table": ( - self._ctx_block_tables - if self._ctx_block_tables is not None - else self._ctx_page_table - ), + # The manager-bound table is the read view rebased per slot above. + "ctx_page_table": ctx_page_table, "block_placement": None, } @@ -2479,6 +2706,7 @@ def _prepare_paged_latent_gen_inputs( accepted_tokens, self._batch_to_slot_buffer.tensor[num_contexts : num_contexts + num_gens], self._ctx_len, + self._ctx_base_pages, self._ctx_block_counts, self._ctx_block_tables, attn_metadata._seq_lens_cuda, @@ -2513,6 +2741,15 @@ def _prepare_paged_latent_gen_inputs( prep.store_offsets, prep.store_mask, ) + # The block decode reads each row shifted past the slot's base, as the + # eager path's dflash_rebased_ctx_view; the kernel already placed the + # block and advertised the context in those rebased positions. + gen_rows = torch.arange( + num_contexts, num_contexts + num_gens, dtype=torch.long, device="cuda" + ) + self._ctx_block_tables_read[gen_rows] = dflash_rebased_ctx_rows( + self._ctx_block_tables[gen_rows], prep.rebase_pages + ) return { "lilicorr_anchor": None, "noise_embedding": noise_embedding, @@ -2520,12 +2757,10 @@ def _prepare_paged_latent_gen_inputs( "num_ctx_per_req": prep.num_ctx, "ctx_k_cache": self._ctx_k_buf, "ctx_v_cache": self._ctx_v_buf, - "ctx_cache_batch_idx": torch.arange( - num_contexts, num_contexts + num_gens, dtype=torch.long, device="cuda" - ), + "ctx_cache_batch_idx": gen_rows, "first_prev_tokens": prep.bonus, "ctx_kv_cache": self._ctx_kv_buf, - "ctx_page_table": self._ctx_block_tables, + "ctx_page_table": self._ctx_block_tables_read, "block_placement": DraftBlockPlacement( prep.blk_pages, prep.blk_offsets, prep.blk_seq_lens, prep.blk_kv_bounds ), diff --git a/tensorrt_llm/_torch/speculative/spec_tail_kernels.py b/tensorrt_llm/_torch/speculative/spec_tail_kernels.py index 6705f605a9eb..aa40462f6049 100644 --- a/tensorrt_llm/_torch/speculative/spec_tail_kernels.py +++ b/tensorrt_llm/_torch/speculative/spec_tail_kernels.py @@ -143,20 +143,25 @@ def _decode_block_table_kernel( offs = tl.arange(0, BLOCK) src_row = src_ptr + row.to(tl.int64) * stride_src_row table_row = table_ptr + row.to(tl.int64) * stride_table - count = tl.zeros([BLOCK], dtype=tl.int64) + extent = tl.zeros([BLOCK], dtype=tl.int64) for start in range(0, n_blocks, BLOCK): cols = start + offs mask = cols < n_blocks encoded = tl.load(src_row + cols * stride_src_col, mask=mask, other=-1).to(tl.int64) - count += (encoded >= 0).to(tl.int64) + extent = tl.maximum(extent, tl.where(encoded >= 0, cols.to(tl.int64) + 1, 0)) tl.store(table_row + cols, (tl.maximum(encoded, 0) // divisor).to(tl.int32), mask=mask) - tl.store(counts_ptr + row, tl.sum(count, axis=0)) + tl.store(counts_ptr + row, tl.max(extent, axis=0)) def decode_draft_block_table( src: torch.Tensor, divisor: int, table: torch.Tensor, counts: torch.Tensor ) -> None: - """``counts[i] = (src[i] >= 0).sum()``; ``table[i] = src[i].clamp(min=0) // divisor``. + """``counts[i]`` = one past the last ``j`` with ``src[i, j] >= 0`` (0 when + none); ``table[i] = src[i].clamp(min=0) // divisor``. + + The extent, not the number of backed entries: a sliding-window pool + releases a row's low entries, and positions are indexed from entry 0, so + only the extent bounds the positions a request may address. Args: src: ``[rows, blocks]`` encoded block offsets (any strides). @@ -212,6 +217,7 @@ class PagedLatentPrepare(NamedTuple): blk_offsets: torch.Tensor # [G, block] int64 blk_seq_lens: torch.Tensor # [G] int32, num_ctx + block blk_kv_bounds: torch.Tensor # [G * block] int32, blk_seq_lens per query token + rebase_pages: torch.Tensor # [G] int64, row entries the block decode skips @triton.jit @@ -221,6 +227,7 @@ def _paged_latent_prepare_kernel( stride_accepted, slots_ptr, ctx_len_ptr, + base_pages_ptr, block_counts_ptr, table_ptr, stride_table, @@ -237,6 +244,7 @@ def _paged_latent_prepare_kernel( blk_offsets_ptr, blk_seq_lens_ptr, blk_kv_bounds_ptr, + rebase_ptr, num_gens, first_row, dummy_slot, @@ -289,8 +297,14 @@ def _paged_latent_prepare_kernel( tl.store(ctx_len_ptr + slot, new_len, mask=gm) tl.store(ctx_len_ptr + dummy_slot, tl.zeros([], dtype=tl.int64)) - # Advertised context, capped so the block still fits in the allocation. + # Advertised context, capped so the block still fits in the allocation, + # then rebased past the slot's base pages (dflash_rebased_ctx_view): the + # block decode reads row entries from ``rebase`` on. num_ctx = tl.minimum(new_len, tl.maximum(capacity - BLOCK_SIZE, 0)) + base = tl.load(base_pages_ptr + slot, mask=gm, other=0).to(tl.int64) + rebase = tl.minimum(base, num_ctx // PAGE_SIZE) + num_ctx = tl.maximum(num_ctx - base * PAGE_SIZE, 0) + tl.store(rebase_ptr + g, rebase, mask=gm) tl.store(num_ctx_ptr + g, num_ctx, mask=gm) tl.store( seq_lens_ptr + row, @@ -300,7 +314,8 @@ def _paged_latent_prepare_kernel( # Block-decode operands: where the block's own latents go, and its bounds. bpos = num_ctx[:, None] + j[None, :] - bpage = tl.load(table_row[:, None] + tl.minimum(bpos // PAGE_SIZE, max_blocks - 1), mask=qmask) + bentry = tl.minimum(rebase[:, None] + bpos // PAGE_SIZE, max_blocks - 1) + bpage = tl.load(table_row[:, None] + bentry, mask=qmask) tl.store(blk_pages_ptr + blk, bpage.to(tl.int64), mask=qmask) tl.store(blk_offsets_ptr + blk, bpos % PAGE_SIZE, mask=qmask) seq_len = (num_ctx + BLOCK_SIZE).to(tl.int32) @@ -317,6 +332,7 @@ def dflash_paged_latent_prepare( accepted_tokens: torch.Tensor, slots: torch.Tensor, ctx_len: torch.Tensor, + base_pages: torch.Tensor, block_counts: torch.Tensor, block_tables: torch.Tensor, seq_lens_cuda: torch.Tensor, @@ -340,7 +356,11 @@ def dflash_paged_latent_prepare( accepted_tokens: ``[batch, width]`` accepted tokens, unit inner stride. slots: ``[G]`` context slots. ctx_len: ``[slots]`` int64 context lengths. - block_counts: ``[batch]`` pages held per batch row. + base_pages: ``[slots]`` int64 pages below each slot's first stored + position; ``num_ctx`` and the block placement are rebased past them + and ``rebase_pages`` says by how many row entries. + block_counts: ``[batch]`` extent of the pages held per batch row + (``decode_draft_block_table``). block_tables: ``[batch, max_blocks]`` int32 page tables, unit inner stride. seq_lens_cuda: ``[batch]`` attention sequence lengths. first_row: batch row of generation request 0. @@ -364,6 +384,7 @@ def dflash_paged_latent_prepare( blk_offsets=torch.empty((num_gens, block_size), **i64), blk_seq_lens=torch.empty((num_gens,), dtype=torch.int32, device=device), blk_kv_bounds=torch.empty((num_gens * block_size,), dtype=torch.int32, device=device), + rebase_pages=torch.empty((num_gens,), **i64), ) if num_gens == 0: return out @@ -374,6 +395,7 @@ def dflash_paged_latent_prepare( accepted_tokens.stride(0), slots, ctx_len, + base_pages, block_counts, block_tables, block_tables.stride(0), diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py index fe729ddd3286..e4b4fdf3a0c4 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py @@ -44,7 +44,7 @@ from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests -from tensorrt_llm._torch.speculative.dflash import DFlashWorker +from tensorrt_llm._torch.speculative.dflash import DFlashSpecMetadata, DFlashWorker from tensorrt_llm.bindings import DataType, SamplingConfig from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE from tensorrt_llm.bindings.executor import KvCacheRetentionConfig @@ -1945,10 +1945,10 @@ def resume_and_restore(request_id, current_cache): # reads position 0 with the whole prompt as its (last) chunk, and the # resize keeps the cache's history. pytest.param(True, False, 8192, 8192, 0, (10243, None), id="draft_view"), - # Disaggregated generation init: the drafter starts a fresh slot at - # context length 0, so the mirror holds no page for the transferred - # prompt (10240 tokens), only the draft headroom. - pytest.param(False, True, 0, 0, 0, (3, None), id="disagg_gen_init-target_view"), + # Disaggregated generation init: the mirror holds no drafter K/V for + # the transferred prompt (10240 tokens), so the whole prompt is + # history; the capacity stays absolute, as the drafter's positions. + pytest.param(False, True, 0, 0, 0, (10243, 10240), id="disagg_gen_init-target_view"), pytest.param(True, True, 0, 0, 0, (10243, None), id="disagg_gen_init-draft_view"), ], ) @@ -2386,8 +2386,8 @@ def _build_fused_kv_buffers(self) -> None: def _get_attention_mask_args(self, layer_idx: int) -> tuple: return True, (-1, -1) - def validate_block_attention_windows(self) -> None: - pass + def validate_block_attention_windows(self) -> frozenset: + return frozenset() def project_target_hidden(self, hidden_states: torch.Tensor) -> torch.Tensor: return hidden_states @@ -2405,32 +2405,15 @@ def precompute_context_kv( return k, k.clone() -def test_disagg_generation_drafter_writes_stay_on_its_own_draft_pages( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """A disaggregated generation worker receives a prompt far larger than its - sliding-window draft pool. The mirror must fit, and the drafter's first - generation step -- a fresh slot that stores the accepted tokens' context - K/V from position 0, then appends its block -- must write only pages its - own mirror holds, never one another request owns.""" - if not torch.cuda.is_available(): - pytest.skip("requires CUDA") - torch.cuda.init() - monkeypatch.setattr( - "tensorrt_llm._torch.speculative.dflash.validate_dflash_trtllm_gen_runtime", - lambda **kwargs: None, - ) - num_layers, num_kv_heads, head_dim, hidden = 2, 1, 64, 8 - window, prompt, occupant_prompt, max_draft_len = 256, 65536, 512, 3 +def _disagg_draft_fixture( + spec: DFlashDecodingConfig, prompt: int, occupant_prompt: int, window: int +) -> tuple[KVCacheManagerV2, LlmRequest, LlmRequest]: + """A sliding-window draft pool with room for 8192 tokens (two layers, one + KV head of 64, bf16), an occupant that prefilled ``occupant_prompt`` + tokens on it, and a request in disaggregated generation init whose + transferred prompt is ``prompt`` tokens.""" + num_layers, num_kv_heads, head_dim = 2, 1, 64 bytes_per_token = num_layers * 2 * num_kv_heads * head_dim * 2 - spec = DFlashDecodingConfig( - max_draft_len=max_draft_len, - speculative_model="draft", - attention_backend="TRTLLM", - mask_token_id=0, - ) - # Room for 8192 tokens: the occupant and a windowed mirror fit, a mirror - # of the whole transferred prompt does not. draft = _small_spec_manager( spec, True, @@ -2452,29 +2435,105 @@ def test_disagg_generation_drafter_writes_stay_on_its_own_draft_pages( occupant.context_chunk_size = occupant_prompt request = LlmRequest( request_id=2, - max_new_tokens=8, + max_new_tokens=64, input_tokens=[1] * prompt, sampling_config=SamplingConfig(1), is_streaming=False, - draft_tokens=[1] * max_draft_len, + draft_tokens=[1] * spec.max_draft_len, ) request.state = LlmRequestState.DISAGG_GENERATION_INIT + return draft, occupant, request + + +def _disagg_generation_step(draft: KVCacheManagerV2, request: LlmRequest) -> None: + """One executor generation step on the draft manager: grow the mirror, + then the post-forward update_resources (nothing accepted, one token).""" + batch = ScheduledRequests() + batch.generation_requests.append(request) + draft._prepare_draft_resources(batch) + request.py_num_accepted_draft_tokens = 0 + request.py_rewind_len = 0 + draft.update_resources(batch) + request.add_new_token(1, 0) + + +def test_gen_init_draft_mirror_survives_generation_updates() -> None: + """A disaggregated generation worker's draft mirror for a prompt far + larger than its sliding-window pool fits at generation init and stays + consistent through generation steps, including update_resources, which + advances the history in absolute positions: the history follows the + request and the pages from the prompt end up to the capacity stay + backed.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + spec = DFlashDecodingConfig(max_draft_len=3, speculative_model="draft") + prompt, page = 65536, 32 + draft, occupant, request = _disagg_draft_fixture(spec, prompt, 512, window=256) + try: + for req in (occupant, request): + batch = ScheduledRequests() + batch.context_requests_last_chunk.append(req) + draft._prepare_draft_resources(batch) + cache = draft.kv_cache_map[request.py_request_id] + # The context worker's first token arrives with the transferred KV. + request.add_new_token(1, 0) + request.state = LlmRequestState.GENERATION_IN_PROGRESS + for _ in range(3): + _disagg_generation_step(draft, request) + assert cache.history_length == request.max_beam_num_tokens - 2 + assert set(range(prompt // page, -(-cache.capacity // page))) <= _backed_ordinals(cache) + draft.free_resources(request) + draft.free_resources(occupant) + finally: + draft.shutdown() + + +@pytest.mark.parametrize("slot_loss", [False, True], ids=["first_step", "slot_loss"]) +@pytest.mark.parametrize("reuse_addressable", [False, True], ids=["unpaired", "reuse_addressable"]) +def test_disagg_generation_drafter_writes_stay_on_its_own_draft_pages( + monkeypatch: pytest.MonkeyPatch, reuse_addressable: bool, slot_loss: bool +) -> None: + """A disaggregated generation worker receives a prompt far larger than its + sliding-window draft pool, while another request owns the pool's low + pages. After generation steps (update_resources included), the drafter's + step -- DFlashSpecMetadata.prepare() assigning the request a fresh slot, + the context K/V store of the accepted tokens, the block decode's read and + its block append -- must address only pages the request's own mirror + holds. The slot is seeded on the page holding the request's absolute + position and reads only what it stored, whether or not the pool joins + block reuse, and again after the request loses its slot and comes back.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + monkeypatch.setattr( + "tensorrt_llm._torch.speculative.dflash.validate_dflash_trtllm_gen_runtime", + lambda **kwargs: None, + ) + num_layers, num_kv_heads, head_dim, hidden = 2, 1, 64, 8 + window, prompt, occupant_prompt, max_draft_len = 256, 65536, 512, 3 + spec = DFlashDecodingConfig( + max_draft_len=max_draft_len, + speculative_model="draft", + attention_backend="TRTLLM", + mask_token_id=0, + ) + draft, occupant, request = _disagg_draft_fixture(spec, prompt, occupant_prompt, window) offsets = torch.zeros( (getattr(draft, "num_attention_op_pools", draft.num_pools), 2, 2, draft.max_blocks_per_seq), dtype=torch.int32, device="cuda", ) + tokens_per_req = max_draft_len + 1 try: - batch = ScheduledRequests() - batch.context_requests_last_chunk.append(occupant) - draft._prepare_draft_resources(batch) - batch = ScheduledRequests() - batch.context_requests_last_chunk.append(request) - draft._prepare_draft_resources(batch) + for req in (occupant, request): + batch = ScheduledRequests() + batch.context_requests_last_chunk.append(req) + draft._prepare_draft_resources(batch) + request.add_new_token(1, 0) request.state = LlmRequestState.GENERATION_IN_PROGRESS - batch = ScheduledRequests() - batch.generation_requests.append(request) - draft._prepare_draft_resources(batch) + for _ in range(2): + _disagg_generation_step(draft, request) worker = DFlashWorker(spec, Mapping()) draft_model = _PoolBoundDFlashDraftModel(num_layers, num_kv_heads, head_dim, hidden) @@ -2495,11 +2554,46 @@ def test_disagg_generation_drafter_writes_stay_on_its_own_draft_pages( draft_kv_cache_manager=draft, ) assert worker._ctx_block_tables is not None, "the drafter did not bind the draft pool" + worker._ctx_reuse_addressable = reuse_addressable # V2 hands out per-layer views over one interleaved pool: a table entry # names layer 0's slot of a page, and layer l of that page is slot # entry + l of the same dense view, so view 0 sees every layer's write. assert worker._ctx_block_divisor == draft.kv_factor pool = worker._ctx_kv_buf[0] + page = worker._ctx_page_size + + spec_metadata = DFlashSpecMetadata( + max_num_requests=2, max_draft_len=max_draft_len, max_total_draft_tokens=max_draft_len + ) + spec_metadata._dflash_worker = worker + + def prepare(request_ids: list[int]) -> None: + """DFlashSpecMetadata.prepare() for a batch of generation rows at + their current positions.""" + reqs = {request.py_request_id: request, occupant.py_request_id: occupant} + spec_metadata.request_ids = request_ids + spec_metadata.num_generations = len(request_ids) + spec_metadata.seq_lens = [tokens_per_req] * len(request_ids) + spec_metadata.host_position_ids = torch.tensor( + [ + reqs[rid].max_beam_num_tokens - 1 + offset + for rid in request_ids + for offset in range(tokens_per_req) + ], + dtype=torch.int32, + ) + spec_metadata.prepare() + + prepare([request.py_request_id]) + if slot_loss: + # The request misses a batch (its slot is evicted), runs a step + # elsewhere in the meantime, and comes back on a fresh slot. + prepare([]) + _disagg_generation_step(draft, request) + prepare([request.py_request_id]) + batch = ScheduledRequests() + batch.generation_requests.append(request) + draft._prepare_draft_resources(batch) def own_slots(req: LlmRequest) -> set[int]: draft.copy_batch_block_offsets(offsets, [req.py_request_id], 1, 0, 1) @@ -2516,16 +2610,10 @@ def own_slots(req: LlmRequest) -> set[int]: assert occupant_slots and request_slots assert not occupant_slots & request_slots - # As DFlashSpecMetadata.prepare() does for a new generation request. - slot = worker._assign_slot(request.py_request_id) - worker._rebuild_batch_to_slot([request.py_request_id]) - assert worker._ctx_len_host[slot] == 0 pool.fill_(7.0) before = pool.clone() - - tokens_per_req = max_draft_len + 1 captured = torch.randn(tokens_per_req, hidden, device="cuda") - spec_metadata = SimpleNamespace( + step_metadata = SimpleNamespace( hidden_size=hidden, runtime_draft_len=max_draft_len, captured_hidden_states=None, @@ -2538,30 +2626,37 @@ def own_slots(req: LlmRequest) -> set[int]: accepted_tokens=torch.ones((1, tokens_per_req), dtype=torch.long, device="cuda"), num_accepted_tokens=torch.tensor([tokens_per_req], dtype=torch.int32, device="cuda"), attn_metadata=attn_metadata, - spec_metadata=spec_metadata, + spec_metadata=step_metadata, draft_model=draft_model, total_target_tokens=tokens_per_req, ) - # The block decode then appends the block's own K/V right after the - # context it reads, through the same table row. - num_ctx = int(inputs["num_ctx_per_req"][0]) - assert num_ctx == tokens_per_req - block = worker._compute_block_size - block_kv = torch.full( - (block, num_layers, num_kv_heads, head_dim), 5.0, dtype=torch.bfloat16, device="cuda" - ) - worker._store_context_kv_paged( - block_kv, - block_kv.clone(), - torch.zeros(block, dtype=torch.long, device="cuda"), - torch.arange(num_ctx, num_ctx + block, dtype=torch.long, device="cuda"), - ) torch.cuda.synchronize() - written = set((pool != before).flatten(1).any(dim=1).nonzero().flatten().tolist()) assert written, "the drafter wrote nothing" assert not written & occupant_slots, "the drafter wrote another request's draft pages" assert written <= request_slots + + # The block decode reads num_ctx positions and appends the block right + # after them, both through the table row it is handed. + num_ctx = int(inputs["num_ctx_per_req"][0]) + read_row = inputs["ctx_page_table"][0].tolist() + block = worker._compute_block_size + addressed = { + read_row[pos // page] + layer + for pos in range(num_ctx + block) + for layer in range(num_layers) + } + assert addressed <= request_slots, "the block decode addresses another request's pages" + + # The slot started empty on the page holding the request's absolute + # position, every page below it hidden, and reads exactly the context + # it stored this step. + base = (request.max_beam_num_tokens - 1) // page + slot = worker._req_to_slot[request.py_request_id] + assert worker._ctx_seed_page_aligned + assert worker._ctx_base_host[slot] == base + assert worker._ctx_len_host[slot] == base * page + assert num_ctx == tokens_per_req draft.free_resources(request) draft.free_resources(occupant) finally: diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py index f88ba9687763..d9b18896ce83 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py @@ -23,10 +23,10 @@ def _worker(req_to_slot): _graph_dummy_id_floor=1 << 30, _dummy_slot=DUMMY_SLOT, _batch_to_slot_buffer=DeviceCopyBuffer(torch.zeros(8, dtype=torch.long)), - _write_ctx_len=lambda updates: None, + _write_ctx_len=lambda updates, base_updates=None: None, ) # Generation requests are already assigned in these tests; prepare assigns nothing. - worker._assign_slot = lambda rid: worker._req_to_slot.get(rid) + worker._assign_slot = lambda rid, **kwargs: worker._req_to_slot.get(rid) return worker diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py index 53958cbab57c..c41fa15259ed 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py @@ -47,3 +47,283 @@ def test_eagle3_one_model_still_gets_a_separate_draft_kv_cache(): config = EagleDecodingConfig(max_draft_len=3, speculative_model="/tmp/eagle3-draft") assert config.spec_dec_mode.is_eagle3_one_model() assert should_use_separate_draft_kv_cache(config) is True + + +@pytest.mark.parametrize( + ("bound_to_pool", "cur", "first_pos", "expected"), + [ + # Managed pool: the block table is indexed by target position, so a + # tail that reuse let the target start at first_pos is stored there. + (True, 0, 8192, 8192), + # A continuing chunk keeps the slot's stored length. + (True, 4096, 4096, 4096), + # Private arena: every request numbers from 0. + (False, 0, 8192, 0), + (False, 0, 0, 0), + ], +) +def test_dflash_ctx_store_start_follows_the_pool_numbering( + bound_to_pool: bool, cur: int, first_pos: int, expected: int +) -> None: + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + worker = object.__new__(DFlashWorker) + worker._ctx_block_tables = object() if bound_to_pool else None + assert worker._ctx_store_start(cur, first_pos) == expected + + +@pytest.mark.parametrize( + ("cur", "start", "reuse_addressable", "expected"), + [ + # Unpaired pool, tail stored at a reused prefix's end: the pages below + # it hold no drafter K/V and are hidden from the block decode. + (0, 8192, False, 128), + # A start inside a page rounds the base up; the tail tokens in that + # page stay stored but unread. + (0, 8200, False, 129), + # A continuing chunk numbers from the stored length: no base change. + (4096, 4096, False, None), + # A pool paired with the target's reuse holds the prefix's drafter K/V. + (0, 8192, True, None), + ], +) +def test_dflash_ctx_base_for_store( + cur: int, start: int, reuse_addressable: bool, expected: int | None +) -> None: + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + worker = object.__new__(DFlashWorker) + worker._ctx_page_size = 64 + worker._ctx_reuse_addressable = reuse_addressable + assert worker._ctx_base_for_store(cur, start) == expected + + +def test_dflash_rebased_ctx_rows_hide_the_reused_prefix() -> None: + """A tail stored after a reused prefix on an unpaired pool. The block decode + reads the slot's row rebased to the tail and a context length shortened by + the prefix, so the never-written prefix pages (poisoned here) are not within + reach, and every tail position still resolves to the page and offset the + store wrote it to.""" + import torch + + from tensorrt_llm._torch.speculative.dflash import dflash_rebased_ctx_rows + + page_size, first_pos, tail = 4, 24, 5 + row = torch.tensor([[7, 3, 11, 0, 5, 9, 13, 2, 6, 1, 15, 4, 8, 10, 12, 14]], dtype=torch.int32) + pool = torch.zeros(16, page_size) + pool[row[0, : first_pos // page_size].long()] = float("nan") + for pos in range(first_pos, first_pos + tail): + pool[row[0, pos // page_size], pos % page_size] = pos + + base_pages = torch.tensor([first_pos // page_size]) + view = dflash_rebased_ctx_rows(row, base_pages) + num_ctx = first_pos + tail - base_pages.item() * page_size + assert view.shape == row.shape + readable = view[0, : -(-num_ctx // page_size)].long() + assert torch.isfinite(pool[readable]).all() + for pos in range(num_ctx): + assert pool[view[0, pos // page_size], pos % page_size] == first_pos + pos + # A slot without a base reads its row unchanged. + assert torch.equal(dflash_rebased_ctx_rows(row, torch.zeros(1, dtype=torch.long)), row) + + +@pytest.mark.parametrize( + ("first_pos", "tail"), + [ + # A reused prefix ending inside a 64-token page with a one-token tail: + # the base rounds up to the next page and the context stays shorter + # than it for 55 steps. + (8200, 1), + # One token into a page: the first write offsets overlap the stored + # tail, which sits below the base and is never read. + (8193, 1), + # A page-aligned prefix end is never short of its base. + (8192, 1), + ], +) +def test_dflash_rebased_ctx_view_writes_the_block_into_an_owned_page( + first_pos: int, tail: int +) -> None: + """A tail stored after a reused prefix on an unpaired pool, with the mirror + funded the way the manager funds it: pages up to the one holding the + context plus the draft tokens, the window below it, nothing past it. The + block decode writes the block's own K/V at rebased positions + num_ctx .. num_ctx + block_size through the view it reads, so at every + context length those positions must resolve to a funded entry, including + the steps on which the context is still shorter than the rounded-up base + and the decode reads nothing. Entry ``base`` is unfunded on those steps and + holds 0, the pool page another request owns.""" + import torch + + from tensorrt_llm._torch.speculative.dflash import DFlashWorker, dflash_rebased_ctx_view + + page_size, block_size, draft_len, window_pages = 64, 3, 3, 64 + worker = object.__new__(DFlashWorker) + worker._ctx_page_size = page_size + worker._ctx_reuse_addressable = False + base = worker._ctx_base_for_store(0, first_pos) + assert base == -(-first_pos // page_size) + base_pages = torch.tensor([base]) + + def funded_row(ctx_len: int) -> torch.Tensor: + last = (ctx_len + draft_len - 1) // page_size + row = torch.zeros(1, base + 8, dtype=torch.int32) + row[0, last + 1 - window_pages : last + 1] = ( + torch.arange(last + 1 - window_pages, last + 1, dtype=torch.int32) + 1000 + ) + return row + + blind_steps = 0 + for ctx_len in range(first_pos + tail, (base + 1) * page_size): + row = funded_row(ctx_len) + view, num_ctx = dflash_rebased_ctx_view(row, base_pages, torch.tensor([ctx_len]), page_size) + assert view.shape == row.shape + blind = ctx_len < base * page_size + blind_steps += blind + assert num_ctx.item() == max(0, ctx_len - base * page_size) + if blind: + assert view[0, 0] == row[0, ctx_len // page_size] + if (ctx_len + draft_len - 1) // page_size < base: + # Entry ``base`` is funded only once the draft tokens reach it. + assert row[0, base] == 0 + for pos in range(num_ctx.item() + block_size): + entry = view[0, pos // page_size] + assert entry >= 1000, (ctx_len, pos) + if pos < num_ctx.item(): + # A read position p is absolute position base * page_size + p. + assert entry == row[0, (base * page_size + pos) // page_size] + assert blind_steps == max(0, base * page_size - first_pos - tail) + + +def test_dflash_write_ctx_len_zero_length_reset_wins_over_a_pending_base() -> None: + """_store_prefill_context records a chunk's base pages before the overflow + check that may reset the same slot to length 0 in the same call. The reset + takes precedence: the slot ends with length 0 and base 0 on the host + mirrors and the device tensors alike, while a slot that was not reset keeps + its new length and base.""" + import torch + + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + worker = object.__new__(DFlashWorker) + worker._ctx_len_host = [0, 512, 64] + worker._ctx_base_host = [0, 8, 1] + worker._ctx_len = torch.tensor(worker._ctx_len_host) + worker._ctx_base_pages = torch.tensor(worker._ctx_base_host) + + # Slot 1: base recorded, then reset by the overflow check. Slot 2: a + # continuing chunk with a new base. + worker._write_ctx_len({1: 0, 2: 96}, {1: 128, 2: 2}) + + assert worker._ctx_len_host == [0, 0, 96] + assert worker._ctx_base_host == [0, 0, 2] + assert worker._ctx_len.tolist() == [0, 0, 96] + assert worker._ctx_base_pages.tolist() == [0, 0, 2] + + +@pytest.mark.parametrize( + ("bound_to_pool", "page_aligned", "reuse_addressable", "position", "expected"), + [ + # A disaggregated generation row: the slot starts empty at the first + # position of the page holding its position and reads from there, + # paired with reuse or not (the mirror of such a request never looked + # a prefix up). + (True, True, False, 8200, (8192, 128)), + (True, True, True, 8200, (8192, 128)), + (True, True, False, 8192, (8192, 128)), + # Clamped to the drafter's context ceiling, base included. + (True, True, False, 1 << 20, (65536, 1024)), + # A pool window shorter than a page may already have released that + # page: the slot starts at the position, reading from the next page. + (True, False, False, 8200, (8200, 129)), + # The private arena is slot-addressed and starts at 0. + (False, True, False, 8200, (0, None)), + ], +) +def test_dflash_generation_slot_seed_is_absolute_with_a_base( + bound_to_pool: bool, + page_aligned: bool, + reuse_addressable: bool, + position: int, + expected: tuple, +) -> None: + """A generation row that gets a fresh slot (a disaggregated generation + worker's first step, or a request whose slot was evicted) starts in the + coordinates of its draft mirror, at a page boundary no further than a page + below its absolute position, with every page below the base hidden from + the block decode.""" + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + worker = object.__new__(DFlashWorker) + worker._ctx_block_tables = object() if bound_to_pool else None + worker._ctx_reuse_addressable = reuse_addressable + worker._ctx_seed_page_aligned = page_aligned + worker._ctx_page_size = 64 + worker._max_ctx = 65536 + updates, bases = {}, {} + worker._seed_generation_slot(3, position, updates, bases) + assert (updates[3], bases.get(3)) == expected + + +def test_dflash_assign_slot_resets_the_base_of_a_reused_slot() -> None: + """A slot freed and handed out again inside one batched update starts + with base 0, not the previous occupant's base.""" + from collections import deque + + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + worker = object.__new__(DFlashWorker) + worker._req_to_slot = {7: 0} + worker._free_slots = deque() + updates, bases = {}, {} + assert worker._assign_slot(8, reset=True, updates=updates, bases=bases) is None + assert worker._assign_slot(7, reset=True, updates=updates, bases=bases) == 0 + assert (updates, bases) == ({0: 0}, {0: 0}) + + +def test_dflash_rollback_restores_the_base_pages() -> None: + """A failed or warmup forward rolls the slot lengths back to the + pre-forward snapshot, and the base pages with them, on the device and the + host mirror; a base seeded during the forward does not survive.""" + import torch + + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + worker = object.__new__(DFlashWorker) + worker._ctx_len_host = [0, 512, 64] + worker._ctx_base_host = [0, 8, 0] + worker._ctx_len = torch.tensor(worker._ctx_len_host) + worker._ctx_base_pages = torch.tensor(worker._ctx_base_host) + worker._req_ctx_pos = {} + worker._saved_ctx_len = worker._ctx_len.clone() + worker._saved_ctx_len_host = list(worker._ctx_len_host) + worker._saved_ctx_base = worker._ctx_base_pages.clone() + worker._saved_ctx_base_host = list(worker._ctx_base_host) + worker._saved_req_ctx_pos = {} + + worker._write_ctx_len({1: 0, 2: 8200}, {2: 129}) + worker._restore_ctx_len_device() + worker._restore_ctx_len_host() + + assert worker._ctx_len.tolist() == [0, 512, 64] + assert worker._ctx_base_pages.tolist() == [0, 8, 0] + assert worker._ctx_len_host == [0, 512, 64] + assert worker._ctx_base_host == [0, 8, 0] + + +@pytest.mark.parametrize( + ("windows", "expected"), + [ + (None, True), + ([None, None], True), + ([None, 4096], True), + ([128, 4096], True), + ([None, 64], False), + ], +) +def test_dflash_windows_span_a_page(windows: list | None, expected: bool) -> None: + """A full-attention layer (window None) keeps every page; only a window + shorter than a page rules out the page-aligned generation seed.""" + from tensorrt_llm._torch.speculative.dflash import dflash_windows_span_a_page + + assert dflash_windows_span_a_page(windows, 128) is expected diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py index 269e1de42197..94796ba0fa13 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py @@ -48,6 +48,8 @@ def _worker(lengths): _max_ctx=1000, _ctx_len=torch.tensor(lengths, dtype=torch.long, device="cuda"), _ctx_len_host=list(lengths), + _ctx_base_pages=torch.zeros(len(lengths), dtype=torch.long, device="cuda"), + _ctx_base_host=[0] * len(lengths), _batch_to_slot_buffer=DeviceCopyBuffer(torch.zeros(8, dtype=torch.long, device="cuda")), ) w._write_ctx_len = partial(DFlashWorker._write_ctx_len, w) @@ -109,9 +111,9 @@ def test_prepare_resets_only_evicted_slots(): written = [] write = worker._write_ctx_len - def record(updates): + def record(updates, base_updates=None): written.append(dict(updates)) - write(updates) + write(updates, base_updates) worker._write_ctx_len = record _prepare(worker, [11, 12, 13, FLOOR + 1], num_generations=4) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_guided_rollback.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_guided_rollback.py index 2e29fd57a5cc..7960f49ba0cd 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_guided_rollback.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_guided_rollback.py @@ -93,6 +93,8 @@ def test_dflash_worker_rolls_back_after_verification(monkeypatch): worker._ctx_len = torch.zeros(batch_size, dtype=torch.long, device="cuda") # _lazy_init_ctx_buffers (stubbed here) normally builds the host mirror. worker._ctx_len_host = [0] * batch_size + worker._ctx_base_pages = torch.zeros(batch_size, dtype=torch.long, device="cuda") + worker._ctx_base_host = [0] * batch_size accepted = torch.zeros(batch_size, K + 1, dtype=torch.int32, device="cuda") # Distinct per-row counts: one row accepts the whole draft, the other only # the golden token. Uniform values would let a broken forward that drops or diff --git a/tests/unittest/_torch/speculative/test_spec_tail_kernels.py b/tests/unittest/_torch/speculative/test_spec_tail_kernels.py index ffddda1808bc..2c11c160187c 100644 --- a/tests/unittest/_torch/speculative/test_spec_tail_kernels.py +++ b/tests/unittest/_torch/speculative/test_spec_tail_kernels.py @@ -64,7 +64,11 @@ def test_strict_accept_matches_eager(num_ctx, num_gens, draft_len): assert torch.equal(accepted[num_ctx:], ref_accepted[num_ctx:]) -def test_decode_draft_block_table_matches_eager(): +def test_decode_draft_block_table_matches_eager() -> None: + """The decoded table is the clamped, divided offsets; the count is each + row's extent, one past its last backed entry, so a sliding window that + released the low entries leaves the count covering the positions the row + still addresses.""" num_pools, rows, n_blocks, divisor = 2, 9, 3000, 6 gen = torch.Generator(device="cuda").manual_seed(1) src = torch.randint( @@ -77,12 +81,18 @@ def test_decode_draft_block_table_matches_eager(): view = src[1, :rows, 0] view[0] = -torch.randint(1, 1 << 20, (n_blocks,), device="cuda", generator=gen).to(torch.int32) view[1] = 0 + # A windowed row: entries 0..99 released, 100..139 backed, the rest unfunded. + view[2] = -1 + view[2, 100:140] = torch.arange(100, 140, device="cuda", dtype=torch.int32) * divisor assert bool((view < 0).float().mean() > 0.3) and bool((view == 0).float().mean() > 0.2) table = torch.full((rows + 3, n_blocks), -7, dtype=torch.int32, device="cuda") counts = torch.full((rows + 3,), -7, dtype=torch.long, device="cuda") decode_draft_block_table(view, divisor, table, counts) encoded = view.to(torch.int64) - assert torch.equal(counts[:rows], (encoded >= 0).sum(dim=1)) + ordinals = torch.arange(1, n_blocks + 1, device="cuda") + extent = torch.where(encoded >= 0, ordinals, 0).max(dim=1).values + assert torch.equal(counts[:rows], extent) + assert counts[2].item() == 140 expected = encoded.clamp(min=0).div(divisor, rounding_mode="floor").to(torch.int32) assert torch.equal(table[:rows], expected) assert counts[0].item() == 0 and bool((table[0] == 0).all()) @@ -143,12 +153,18 @@ def _paged_worker(draft_len, page_size, max_blocks, num_slots, layers, head_dim, perm = torch.randperm(pages, device="cuda", generator=gen).to(torch.int32) worker._ctx_block_tables = perm[: num_slots * max_blocks].view(num_slots, max_blocks).clone() worker._ctx_block_counts = torch.zeros(num_slots, dtype=torch.long, device="cuda") + worker._ctx_base_pages = torch.zeros(num_slots, dtype=torch.long, device="cuda") + worker._ctx_block_tables_read = torch.zeros_like(worker._ctx_block_tables) worker._resolved_mask_token_id = 3 return worker, storage +@pytest.mark.parametrize("based", [False, True], ids=["no_base", "based"]) @pytest.mark.parametrize("draft_len", [3, 7]) -def test_paged_latent_prepare_matches_eager(draft_len): +def test_paged_latent_prepare_matches_eager(draft_len: int, based: bool) -> None: + """The fused generation-step preparation matches the eager one, including + the read view rebased past each slot's base pages: the advertised context, + the rebased block-table rows and the block placement.""" vocab, hidden, features, layers, head_dim = 100, 64, 128, 3, 96 page_size, max_blocks, num_slots, max_ctx = 16, 6, 10, 70 num_ctx_reqs, num_gens = 1, 6 @@ -166,6 +182,11 @@ def test_paged_latent_prepare_matches_eager(draft_len): worker._batch_to_slot_buffer.tensor[:batch] = slots worker._ctx_len[:5] = torch.tensor([10, 3, 66, 5, 40], device="cuda") worker._ctx_block_counts[:batch] = torch.tensor([0, 3, 6, 2, 0, 1, 1], device="cuda") + if based: + # Slot 2 (row 2) advertises past its base, so its read row shifts by + # it; slot 4 (row 3) advertises less than its base, reads nothing and + # shifts only to the page holding its context. The dummy slot keeps 0. + worker._ctx_base_pages[:5] = torch.tensor([0, 1, 1, 1, 2], device="cuda") num_accepted = torch.tensor([1, 2, width, 3, 1, 2, 1], dtype=torch.int32, device="cuda") accepted = torch.randint( 0, vocab, (batch, width), dtype=torch.int32, device="cuda", generator=gen @@ -193,6 +214,7 @@ def attn_metadata(): def prepare(fused): w = copy.copy(worker) w._ctx_len = worker._ctx_len.clone() + w._ctx_block_tables_read = worker._ctx_block_tables_read.clone() pool_storage = storage.clone() w._ctx_kv_buf = [pool_storage[:, layer] for layer in range(layers)] w._ctx_kv_ptrs = torch.tensor([p.data_ptr() for p in w._ctx_kv_buf], device="cuda") @@ -253,9 +275,13 @@ def prepare(fused): assert bool(shared.any()) and bool(untouched.any()) rows = eager["ctx_cache_batch_idx"] + read_rows = eager["ctx_page_table"].index_select(0, rows) + assert torch.equal(fused["ctx_page_table"].index_select(0, rows), read_rows) + if based: + assert not torch.equal(read_rows, worker._ctx_block_tables.index_select(0, rows)) fixup = _build_mla_block_fixup( eager["num_ctx_per_req"], - worker._ctx_block_tables.index_select(0, rows).long(), + read_rows.long(), draft_len, page_size, ) From 89157212a6d2151d1d0ae15fe50063013091ce64 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Fri, 30 Oct 2026 12:39:04 -0700 Subject: [PATCH 04/14] [TRTLLM-17390][fix] Store a reused tail at its target position on a paired draft pool A draft pool paired with the target's block reuse holds the prefix's drafter K/V below the tail, so the tail is stored at its target position whether or not the worker reads the rebased view; the Kimi K3 workers, which read from position 0, no longer overwrite the shared prefix pages. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/speculative/dflash.py | 10 ++++--- .../hw_agnostic/test_dflash_draft_kv.py | 26 ++++++++++++++----- 2 files changed, 27 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 733da12b36a1..1a0417652899 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -1561,14 +1561,18 @@ def _ctx_store_start(self, cur: int, first_pos: int) -> int: the tail reaches the capture buffer. On the managed pool the block table is indexed by target position, so - the tail is stored at ``first_pos``. The private arena numbers every - request from 0, where the tail is all the context there is. + the tail is stored at ``first_pos``: always on a pool paired with the + target's reuse, which holds the prefix's drafter K/V below it, and on + an unpaired pool when the worker reads the rebased view, which hides + the empty prefix. Otherwise (the private arena, or an unpaired pool + read from position 0) every request numbers from 0, where the tail is + all the context there is. """ if ( cur == 0 and first_pos > 0 and self._ctx_block_tables is not None - and self._ctx_reads_rebased_view + and (self._ctx_reuse_addressable or self._ctx_reads_rebased_view) ): return first_pos return cur diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py index c41fa15259ed..d89b92753a90 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py @@ -50,25 +50,39 @@ def test_eagle3_one_model_still_gets_a_separate_draft_kv_cache(): @pytest.mark.parametrize( - ("bound_to_pool", "cur", "first_pos", "expected"), + ("bound_to_pool", "reuse_addressable", "reads_rebased_view", "cur", "first_pos", "expected"), [ # Managed pool: the block table is indexed by target position, so a # tail that reuse let the target start at first_pos is stored there. - (True, 0, 8192, 8192), + (True, False, True, 0, 8192, 8192), + (True, True, True, 0, 8192, 8192), + # A pool paired with the target's reuse holds the prefix below the + # tail, so the tail goes on top of it even for a worker that reads + # from position 0 (the Kimi K3 workers). + (True, True, False, 0, 8192, 8192), + # An unpaired pool read from position 0 numbers the tail from 0. + (True, False, False, 0, 8192, 0), # A continuing chunk keeps the slot's stored length. - (True, 4096, 4096, 4096), + (True, False, True, 4096, 4096, 4096), # Private arena: every request numbers from 0. - (False, 0, 8192, 0), - (False, 0, 0, 0), + (False, False, True, 0, 8192, 0), + (False, False, True, 0, 0, 0), ], ) def test_dflash_ctx_store_start_follows_the_pool_numbering( - bound_to_pool: bool, cur: int, first_pos: int, expected: int + bound_to_pool: bool, + reuse_addressable: bool, + reads_rebased_view: bool, + cur: int, + first_pos: int, + expected: int, ) -> None: from tensorrt_llm._torch.speculative.dflash import DFlashWorker worker = object.__new__(DFlashWorker) worker._ctx_block_tables = object() if bound_to_pool else None + worker._ctx_reuse_addressable = reuse_addressable + worker._ctx_reads_rebased_view = reads_rebased_view assert worker._ctx_store_start(cur, first_pos) == expected From 277eb55320ef430a6386b8db7a5e3ad34e6d9ce1 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Fri, 30 Oct 2026 13:30:38 -0700 Subject: [PATCH 05/14] [TRTLLM-17390][fix] Keep a page below the windows of a DFlash/DSpark draft pool A drafter slot seeded on the page holding the request's position trails the request by up to a page less one token, so its window reached entries the mirror had released. DFlash/DSpark draft pools retain one more page per window, in the pool, the quota conversion and the budget-split estimators. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- .../kv_cache/kv_cache_manager_v2.py | 44 ++++++- tensorrt_llm/_torch/speculative/dflash.py | 11 +- .../kv_cache/sharing/test_staging_lender.py | 7 +- .../kv_cache/test_kvcm2_integration.py | 114 +++++++++++++++++- 4 files changed, 164 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index dd234f3bf8ce..49ab2489925c 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -763,6 +763,24 @@ def _extend_swa_windows_for_reuse( return retention_windows +def _draft_slot_trail_tokens( + spec_config: Optional["DecodingBaseConfig"], *, is_draft: bool, tokens_per_block: int +) -> int: + """SWA retention a DFlash/DSpark draft pool keeps below its windows. + + The drafter seeds a slot that holds no context yet at the first position of + the page holding the request's position (``DFlashWorker._seed_generation_slot``), + so the slot's positions trail the request's, which the mirror's history + follows, by up to a page less one token. The far edge of the drafter's + window reaches that far below the mirror's window; one more page of + retention keeps every entry it reads on a page the request holds. + """ + if not is_draft or spec_config is None: + return 0 + mode = spec_config.spec_dec_mode + return int(tokens_per_block) if mode.is_dflash() or mode.is_dspark() else 0 + + def _derive_request_windows_from_layer_types( model_config: ModelConfigPython, local_layer_ids: Sequence[int], @@ -2127,6 +2145,15 @@ def __init__( self.reuse_match_backoff, max_seq_len, ) + # The drafter's own windows come from its model config, so this widens + # only what the pool retains (see _draft_slot_trail_tokens). + self.max_attention_window_vec = _extend_swa_windows_for_reuse( + self.max_attention_window_vec, + _draft_slot_trail_tokens( + spec_config, is_draft=is_draft, tokens_per_block=tokens_per_block + ), + max_seq_len, + ) event_window_size = max( self.max_seq_len if window_size is None else int(window_size) @@ -7686,6 +7713,13 @@ def get_cache_size_per_token( backoff, max_seq_len, ) + attention_windows = _extend_swa_windows_for_reuse( + attention_windows, + _draft_slot_trail_tokens( + spec_config, is_draft=is_draft, tokens_per_block=tokens_per_block + ), + max_seq_len, + ) _, generation_capacity_headroom = _get_generation_kv_capacity( spec_config, is_draft=is_draft ) @@ -7826,7 +7860,8 @@ def _get_per_request_layer_tokens( return None # ``covered_windows`` are the windows ``get_cache_size_per_token`` # charges as fixed per-request retention in the affine intercept; keep - # them in lockstep with its accounting (reuse backoff included). + # them in lockstep with its accounting (reuse backoff and the drafter's + # trailing page included). if ( kv_cache_config is not None and kv_cache_config.enable_block_reuse @@ -7836,6 +7871,13 @@ def _get_per_request_layer_tokens( backoff = draft_prompt_lookahead(spec_config) or 0 covered_windows = _extend_swa_windows_for_reuse(covered_windows, backoff, max_seq_len) + covered_windows = _extend_swa_windows_for_reuse( + covered_windows, + _draft_slot_trail_tokens( + spec_config, is_draft=is_draft, tokens_per_block=tokens_per_block + ), + max_seq_len, + ) if is_draft: # Draft mirrors are reserved at the request's full context length # (measured: a 41k-token context charges 41k tokens per mirror diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 1a0417652899..4e353171ca93 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -1607,10 +1607,13 @@ def _seed_generation_slot(self, slot: int, position: int, updates: dict, bases: and the block decode reads from its base, never below it: * at the first position of the page holding ``position``. The slot's - positions then trail the absolute ones by less than a page, so they - stay on pages the mirror backs as long as every window of the pool - spans a page (``_ctx_seed_page_aligned``), and the slot reads what - it stores from its first step; + positions then trail the absolute ones by less than a page, and the + mirror keeps one page below each window for that trail + (``_draft_slot_trail_tokens`` in kv_cache_manager_v2), so every + entry the window reads stays on a page the request holds. Taken + when every window of the pool spans a page + (``_ctx_seed_page_aligned``); the slot reads what it stores from + its first step; * otherwise at ``position`` itself, with the base rounded up to the next page: the slot reads nothing until its context reaches it. diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/sharing/test_staging_lender.py b/tests/unittest/_torch/pyexecutor/kv_cache/sharing/test_staging_lender.py index 79717b94f868..b8b5a3deaa12 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/sharing/test_staging_lender.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/sharing/test_staging_lender.py @@ -2704,9 +2704,10 @@ def test_a_windowed_joint_draft_pool_resumes_only_where_its_context_resize_holds ) as draft, ): pools = (target, draft) - for mgr in pools: - assert JOINT_WINDOW in kit.windows(mgr), "a pool has no window: nothing checked" - beg, end = kit.stale_blocks(mgr, kit.windows(mgr).index(JOINT_WINDOW), PAIRED_END) + # A DFlash draft pool retains one page below its window (_draft_slot_trail_tokens). + for mgr, window in ((target, JOINT_WINDOW), (draft, JOINT_WINDOW + TPB)): + assert window in kit.windows(mgr), "a pool has no window: nothing checked" + beg, end = kit.stale_blocks(mgr, kit.windows(mgr).index(window), PAIRED_END) assert end <= beg, "a window released blocks at the fetch's end: nothing checked" scheduler = KVCacheV2Scheduler( max_batch_size=4, diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py index e4b4fdf3a0c4..dce0246a1fcc 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py @@ -2219,9 +2219,10 @@ def _small_spec_manager( max_gpu_total_bytes: int = 256 << 20, max_seq_len: int = 16384, max_num_tokens: int = 8192, + tokens_per_block: int = 32, ) -> KVCacheManagerV2: - """A two-layer, head_dim 64, 32-token-page V2 manager for ``spec``, - optionally with a uniform sliding window and no block reuse.""" + """A two-layer, head_dim 64 V2 manager for ``spec`` (32-token pages by + default), optionally with a uniform sliding window and no block reuse.""" num_layers = 2 config = KvCacheConfig( enable_block_reuse=False, @@ -2234,7 +2235,7 @@ def _small_spec_manager( num_layers=num_layers, num_kv_heads=num_kv_heads, head_dim=64, - tokens_per_block=32, + tokens_per_block=tokens_per_block, max_seq_len=max_seq_len, max_batch_size=2, max_num_tokens=max_num_tokens, @@ -2406,7 +2407,11 @@ def precompute_context_kv( def _disagg_draft_fixture( - spec: DFlashDecodingConfig, prompt: int, occupant_prompt: int, window: int + spec: DFlashDecodingConfig, + prompt: int, + occupant_prompt: int, + window: int, + tokens_per_block: int = 32, ) -> tuple[KVCacheManagerV2, LlmRequest, LlmRequest]: """A sliding-window draft pool with room for 8192 tokens (two layers, one KV head of 64, bf16), an occupant that prefilled ``occupant_prompt`` @@ -2423,6 +2428,7 @@ def _disagg_draft_fixture( max_gpu_total_bytes=8192 * bytes_per_token, max_seq_len=1 << 17, max_num_tokens=occupant_prompt, + tokens_per_block=tokens_per_block, ) occupant = LlmRequest( request_id=1, @@ -2663,6 +2669,106 @@ def own_slots(req: LlmRequest) -> set[int]: draft.shutdown() +@pytest.mark.parametrize("lag", [2, 7, 31, 127]) +def test_trailing_drafter_slot_reads_only_pages_its_request_holds( + monkeypatch: pytest.MonkeyPatch, lag: int +) -> None: + """A disaggregated generation worker seeds the drafter slot on the page + holding the request's position, ``lag`` tokens behind it, while the + request's sliding-window draft mirror advances its history with the + request. Over many generation steps (one accepted token each), every + block-table entry the drafter's window reads -- keys from the slot length + + 1 - window, never below the slot's base -- must stay on a page the + request's own mirror holds; a released entry resolves to another + request's page.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + monkeypatch.setattr( + "tensorrt_llm._torch.speculative.dflash.validate_dflash_trtllm_gen_runtime", + lambda **kwargs: None, + ) + num_layers, num_kv_heads, head_dim, hidden = 2, 1, 64, 8 + window, page, max_draft_len, steps = 256, 128, 3, 600 + spec = DFlashDecodingConfig( + max_draft_len=max_draft_len, + speculative_model="draft", + attention_backend="TRTLLM", + mask_token_id=0, + ) + # The first token and two generation steps put the request at prompt + 2. + prompt = 65536 + lag - 2 + draft, occupant, request = _disagg_draft_fixture(spec, prompt, 512, window, page) + offsets = torch.zeros( + (getattr(draft, "num_attention_op_pools", draft.num_pools), 2, 2, draft.max_blocks_per_seq), + dtype=torch.int32, + device="cuda", + ) + try: + for req in (occupant, request): + batch = ScheduledRequests() + batch.context_requests_last_chunk.append(req) + draft._prepare_draft_resources(batch) + request.add_new_token(1, 0) + request.state = LlmRequestState.GENERATION_IN_PROGRESS + for _ in range(2): + _disagg_generation_step(draft, request) + + worker = DFlashWorker(spec, Mapping()) + draft_model = _PoolBoundDFlashDraftModel(num_layers, num_kv_heads, head_dim, hidden) + worker.set_draft_model(draft_model) + attn_metadata = SimpleNamespace( + max_seq_len=1 << 17, + draft_kv_cache_block_offsets=offsets, + num_contexts=0, + num_seqs=1, + num_ctx_tokens=0, + _seq_lens_cuda=torch.zeros(1, dtype=torch.int32, device="cuda"), + _seq_lens=torch.zeros(1, dtype=torch.int32), + ) + worker._lazy_init_ctx_buffers( + draft_model, + SimpleNamespace(max_num_requests=2), + attn_metadata, + draft_kv_cache_manager=draft, + ) + assert worker._ctx_page_size == page + spec_metadata = DFlashSpecMetadata( + max_num_requests=2, max_draft_len=max_draft_len, max_total_draft_tokens=max_draft_len + ) + spec_metadata._dflash_worker = worker + spec_metadata.request_ids = [request.py_request_id] + spec_metadata.num_generations = 1 + spec_metadata.seq_lens = [max_draft_len + 1] + position = request.max_beam_num_tokens - 1 + spec_metadata.host_position_ids = torch.tensor( + [position + offset for offset in range(max_draft_len + 1)], dtype=torch.int32 + ) + spec_metadata.prepare() + slot = worker._req_to_slot[request.py_request_id] + seed_len, base = int(worker._ctx_len_host[slot]), int(worker._ctx_base_host[slot]) + assert (base, position - seed_len) == (position // page, lag) + + cache = draft.kv_cache_map[request.py_request_id] + unbacked_reads = [] + for step in range(steps): + slot_len = seed_len + step + 1 + lowest = max(base * page, slot_len + 1 - window) + read = {key // page for key in range(lowest, slot_len)} + missing = read - _backed_ordinals(cache) + if missing: + unbacked_reads.append((step, cache.history_length, slot_len, sorted(missing))) + _disagg_generation_step(draft, request) + assert not unbacked_reads, ( + f"{len(unbacked_reads)}/{steps} steps read a released entry; first " + f"(step, history, slot length, entries): {unbacked_reads[:3]}" + ) + draft.free_resources(request) + draft.free_resources(occupant) + finally: + draft.shutdown() + + def _make_publishing_manager(policy: BlockReusePolicy) -> KVCacheManagerV2: """A manager wired just far enough to run the publish/history bookkeeping.""" manager = object.__new__(KVCacheManagerV2) From 74bb2021c319526607c4346d9535f470d70fea66 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Fri, 30 Oct 2026 13:50:22 -0700 Subject: [PATCH 06/14] [TRTLLM-17390][fix] Seed a drafter slot on a page boundary only over a trailing page DFlashWorker starts a seeded slot on the page holding its position only when every window of the bound pool is backed below by the trail page the V2 draft manager keeps (draft_slot_trail_tokens); a windowed V1 pool takes the position-aligned seed. Annotate the dummy-slot test's recorder. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py | 9 ++++----- tensorrt_llm/_torch/speculative/dflash.py | 10 ++++++++-- .../pyexecutor/kv_cache/test_kvcm2_integration.py | 1 + .../speculative/hw_agnostic/test_dflash_dummy_slot.py | 2 +- 4 files changed, 14 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 49ab2489925c..6cbfab84d9b3 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -2147,12 +2147,11 @@ def __init__( ) # The drafter's own windows come from its model config, so this widens # only what the pool retains (see _draft_slot_trail_tokens). + self.draft_slot_trail_tokens = _draft_slot_trail_tokens( + spec_config, is_draft=is_draft, tokens_per_block=tokens_per_block + ) self.max_attention_window_vec = _extend_swa_windows_for_reuse( - self.max_attention_window_vec, - _draft_slot_trail_tokens( - spec_config, is_draft=is_draft, tokens_per_block=tokens_per_block - ), - max_seq_len, + self.max_attention_window_vec, self.draft_slot_trail_tokens, max_seq_len ) event_window_size = max( diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 4e353171ca93..afba85439a7e 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -1497,8 +1497,14 @@ def _lazy_init_ctx_buffers( self._ctx_reuse_addressable = bool( getattr(draft_kv_cache_manager, "enable_joint_kv_cache_reuse", False) ) - self._ctx_seed_page_aligned = dflash_windows_span_a_page( - getattr(draft_kv_cache_manager, "max_attention_window_vec", None), page_size + # A slot seeded on a page boundary trails its request by less + # than a page; a windowed pool backs that trail only when it + # keeps a page below its windows (kv_cache_manager_v2's + # draft_slot_trail_tokens; a V1 pool does not). + windows = getattr(draft_kv_cache_manager, "max_attention_window_vec", None) + self._ctx_seed_page_aligned = dflash_windows_span_a_page(windows, page_size) and ( + all(window is None for window in windows or ()) + or getattr(draft_kv_cache_manager, "draft_slot_trail_tokens", 0) >= page_size ) self._ctx_kv_last_page_len = torch.full( (num_slots,), page_size, dtype=torch.int32, device="cuda" diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py index dce0246a1fcc..f7bf22d766a1 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py @@ -2733,6 +2733,7 @@ def test_trailing_drafter_slot_reads_only_pages_its_request_holds( draft_kv_cache_manager=draft, ) assert worker._ctx_page_size == page + assert worker._ctx_seed_page_aligned spec_metadata = DFlashSpecMetadata( max_num_requests=2, max_draft_len=max_draft_len, max_total_draft_tokens=max_draft_len ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py index 94796ba0fa13..7056c17c3ea7 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py @@ -111,7 +111,7 @@ def test_prepare_resets_only_evicted_slots(): written = [] write = worker._write_ctx_len - def record(updates, base_updates=None): + def record(updates: dict[int, int], base_updates: dict[int, int] | None = None) -> None: written.append(dict(updates)) write(updates, base_updates) From e8db038d2eb9b80f505259f10f2ca1ea6d6126f0 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Fri, 30 Oct 2026 13:55:19 -0700 Subject: [PATCH 07/14] [TRTLLM-17390][fix] Describe the draft mirror of a reused-prefix chunk as stored The test case comment states the absolute store position and the mirror's kept history. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache/test_kvcm2_integration.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py index f7bf22d766a1..10e7be33f62f 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py @@ -1938,8 +1938,8 @@ def resume_and_restore(request_id, current_cache): # lowers it. pytest.param(False, False, 8192, 8192, 8704, (9216, 8704), id="history_ahead"), # A mirror first seen past position 0 (a target prefix this worker did - # not compute): its drafter slot stores the chunk from 0, so nothing - # below the chunk is history and the mirror covers [0, chunk end). + # not compute) keeps its history and covers [0, chunk end); the DFlash + # drafter stores the chunk at its target position above the prefix. pytest.param(False, False, 8192, 0, 0, (9216, None), id="fresh_past_zero"), # MTP, Eagle and two-engine drafters keep the draft request view: it # reads position 0 with the whole prompt as its (last) chunk, and the From d7009512c7351e0121757e1bb30a7c600d8ed65f Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Fri, 30 Oct 2026 14:22:48 -0700 Subject: [PATCH 08/14] [TRTLLM-17390][fix] Take a reused prefix as history on a fresh DFlash/DSpark mirror A DFlash/DSpark drafter stores a chunk after a prefix it did not compute at its target position over a base hiding the prefix, so its mirror takes the chunk start as history and a sliding window trims the prefix instead of backing it. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- .../kv_cache/kv_cache_manager_v2.py | 13 +++++++--- .../kv_cache/test_kvcm2_integration.py | 24 ++++++++++++------- 2 files changed, 25 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 6cbfab84d9b3..a477f499d2c1 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -6182,14 +6182,21 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): kv_cache.history_length, req.context_current_position + req.context_chunk_size, ) - elif kv_cache.capacity >= req.context_current_position: + elif ( + kv_cache.capacity >= req.context_current_position + or self.draft_slot_trail_tokens > 0 + ): # The chunk start is history for a mirror that already # holds [0, chunk start): sliding-window layers keep # pages only from history + 1 - window (AttnLifeCycle:: # getStaleRange), so the mirror holds its window plus # the chunk in flight rather than the whole prompt so - # far. A mirror first seen past position 0 is not one - # and keeps its history, i.e. pages from 0. + # far. So it is for a DFlash/DSpark mirror first seen + # past position 0 (a prefix this worker did not + # compute): that drafter stores the chunk at its target + # position over a base hiding the prefix + # (DFlashWorker._ctx_store_start). Any other mirror + # first seen there keeps its history, i.e. pages from 0. history = max(kv_cache.history_length, req.context_current_position) if not kv_cache.resize(capacity, history): raise self._mirror_resize_error(req, "context", capacity) diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py index 10e7be33f62f..87d90794b72f 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py @@ -1925,31 +1925,35 @@ def resume_and_restore(request_id, current_cache): "mirror_capacity", "mirror_history", "expected_resize", + "trail_tokens", ), [ # External drafter (target view), aggregated chunked prefill. The # prompt is 8192 + 2 * 1024 tokens, so the 1024-token chunk at 8192 is # not the last one and carries no draft tokens. - pytest.param(False, False, 0, 0, 0, (1024, 0), id="first_chunk"), + pytest.param(False, False, 0, 0, 0, (1024, 0), 0, id="first_chunk"), # A later chunk of a mirror that holds [0, 8192): the chunk start is # history, so a sliding window drops the pages below it. - pytest.param(False, False, 8192, 8192, 4096, (9216, 8192), id="later_chunk"), + pytest.param(False, False, 8192, 8192, 4096, (9216, 8192), 0, id="later_chunk"), # The history is already ahead of the chunk start; a resize never # lowers it. - pytest.param(False, False, 8192, 8192, 8704, (9216, 8704), id="history_ahead"), + pytest.param(False, False, 8192, 8192, 8704, (9216, 8704), 0, id="history_ahead"), # A mirror first seen past position 0 (a target prefix this worker did - # not compute) keeps its history and covers [0, chunk end); the DFlash - # drafter stores the chunk at its target position above the prefix. - pytest.param(False, False, 8192, 0, 0, (9216, None), id="fresh_past_zero"), + # not compute). A DFlash/DSpark pool (one with a trailing page) takes + # the chunk start as history: that drafter stores the chunk at its + # target position above the prefix. Any other keeps its history and + # covers [0, chunk end). + pytest.param(False, False, 8192, 0, 0, (9216, 8192), 32, id="fresh_past_zero-dflash"), + pytest.param(False, False, 8192, 0, 0, (9216, None), 0, id="fresh_past_zero-pard"), # MTP, Eagle and two-engine drafters keep the draft request view: it # reads position 0 with the whole prompt as its (last) chunk, and the # resize keeps the cache's history. - pytest.param(True, False, 8192, 8192, 0, (10243, None), id="draft_view"), + pytest.param(True, False, 8192, 8192, 0, (10243, None), 0, id="draft_view"), # Disaggregated generation init: the mirror holds no drafter K/V for # the transferred prompt (10240 tokens), so the whole prompt is # history; the capacity stays absolute, as the drafter's positions. - pytest.param(False, True, 0, 0, 0, (10243, 10240), id="disagg_gen_init-target_view"), - pytest.param(True, True, 0, 0, 0, (10243, None), id="disagg_gen_init-draft_view"), + pytest.param(False, True, 0, 0, 0, (10243, 10240), 0, id="disagg_gen_init-target_view"), + pytest.param(True, True, 0, 0, 0, (10243, None), 0, id="disagg_gen_init-draft_view"), ], ) def test_unpaired_draft_mirror_context_resize( @@ -1959,6 +1963,7 @@ def test_unpaired_draft_mirror_context_resize( mirror_capacity: int, mirror_history: int, expected_resize: tuple, + trail_tokens: int, ) -> None: """Which request view an unpaired draft mirror reads, and the capacity and history its context resize asks for.""" @@ -1987,6 +1992,7 @@ def test_unpaired_draft_mirror_context_resize( manager.is_draft = True manager.enable_joint_kv_cache_reuse = False manager._mirror_uses_draft_view = uses_draft_view + manager.draft_slot_trail_tokens = trail_tokens manager.num_extra_kv_tokens = 0 manager.kv_cache_map = {request.py_request_id: kv_cache} manager._resume_and_restore = Mock(return_value=True) From 65dde5e323d39fb07cc2b92e863284d93a330b07 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Mon, 2 Nov 2026 13:27:27 -0800 Subject: [PATCH 09/14] [TRTLLM-17390][fix] Check the spec mode for a DFlash/DSpark draft mirror The context resize of an unpaired draft mirror first seen past position 0 takes the chunk start as history when a constructor flag set from the DFlash/DSpark spec mode says the drafter stores at target positions. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- .../kv_cache/kv_cache_manager_v2.py | 12 +++++--- .../kv_cache/test_kvcm2_integration.py | 29 +++++++++---------- 2 files changed, 22 insertions(+), 19 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index a477f499d2c1..b2c87d7f2dc6 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -2080,6 +2080,13 @@ def __init__( self._mirror_uses_draft_view = not ( is_draft and spec_config is not None and spec_config.spec_dec_mode.is_external_drafter() ) + # DFlash and DSpark drafters store a chunk at its target position + # (DFlashWorker._ctx_store_start). + self._is_dflash_draft = ( + is_draft + and spec_config is not None + and (spec_config.spec_dec_mode.is_dflash() or spec_config.spec_dec_mode.is_dspark()) + ) # Mirror V1: expose max_draft_len so the native disagg AuxBuffer # (_make_aux_buffer's getattr fallback) is sized for MTP/spec decoding. self.max_draft_len = spec_config.max_draft_len if spec_config is not None else 0 @@ -6182,10 +6189,7 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): kv_cache.history_length, req.context_current_position + req.context_chunk_size, ) - elif ( - kv_cache.capacity >= req.context_current_position - or self.draft_slot_trail_tokens > 0 - ): + elif kv_cache.capacity >= req.context_current_position or self._is_dflash_draft: # The chunk start is history for a mirror that already # holds [0, chunk start): sliding-window layers keep # pages only from history + 1 - window (AttnLifeCycle:: diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py index 87d90794b72f..ed98c6687995 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py @@ -1925,35 +1925,34 @@ def resume_and_restore(request_id, current_cache): "mirror_capacity", "mirror_history", "expected_resize", - "trail_tokens", + "dflash", ), [ # External drafter (target view), aggregated chunked prefill. The # prompt is 8192 + 2 * 1024 tokens, so the 1024-token chunk at 8192 is # not the last one and carries no draft tokens. - pytest.param(False, False, 0, 0, 0, (1024, 0), 0, id="first_chunk"), + pytest.param(False, False, 0, 0, 0, (1024, 0), False, id="first_chunk"), # A later chunk of a mirror that holds [0, 8192): the chunk start is # history, so a sliding window drops the pages below it. - pytest.param(False, False, 8192, 8192, 4096, (9216, 8192), 0, id="later_chunk"), + pytest.param(False, False, 8192, 8192, 4096, (9216, 8192), False, id="later_chunk"), # The history is already ahead of the chunk start; a resize never # lowers it. - pytest.param(False, False, 8192, 8192, 8704, (9216, 8704), 0, id="history_ahead"), + pytest.param(False, False, 8192, 8192, 8704, (9216, 8704), False, id="history_ahead"), # A mirror first seen past position 0 (a target prefix this worker did - # not compute). A DFlash/DSpark pool (one with a trailing page) takes - # the chunk start as history: that drafter stores the chunk at its - # target position above the prefix. Any other keeps its history and - # covers [0, chunk end). - pytest.param(False, False, 8192, 0, 0, (9216, 8192), 32, id="fresh_past_zero-dflash"), - pytest.param(False, False, 8192, 0, 0, (9216, None), 0, id="fresh_past_zero-pard"), + # not compute). A DFlash/DSpark pool takes the chunk start as history: + # that drafter stores the chunk at its target position above the + # prefix. Any other keeps its history and covers [0, chunk end). + pytest.param(False, False, 8192, 0, 0, (9216, 8192), True, id="fresh_past_zero-dflash"), + pytest.param(False, False, 8192, 0, 0, (9216, None), False, id="fresh_past_zero-pard"), # MTP, Eagle and two-engine drafters keep the draft request view: it # reads position 0 with the whole prompt as its (last) chunk, and the # resize keeps the cache's history. - pytest.param(True, False, 8192, 8192, 0, (10243, None), 0, id="draft_view"), + pytest.param(True, False, 8192, 8192, 0, (10243, None), False, id="draft_view"), # Disaggregated generation init: the mirror holds no drafter K/V for # the transferred prompt (10240 tokens), so the whole prompt is # history; the capacity stays absolute, as the drafter's positions. - pytest.param(False, True, 0, 0, 0, (10243, 10240), 0, id="disagg_gen_init-target_view"), - pytest.param(True, True, 0, 0, 0, (10243, None), 0, id="disagg_gen_init-draft_view"), + pytest.param(False, True, 0, 0, 0, (10243, 10240), False, id="disagg_gen_init-target_view"), + pytest.param(True, True, 0, 0, 0, (10243, None), False, id="disagg_gen_init-draft_view"), ], ) def test_unpaired_draft_mirror_context_resize( @@ -1963,7 +1962,7 @@ def test_unpaired_draft_mirror_context_resize( mirror_capacity: int, mirror_history: int, expected_resize: tuple, - trail_tokens: int, + dflash: bool, ) -> None: """Which request view an unpaired draft mirror reads, and the capacity and history its context resize asks for.""" @@ -1992,7 +1991,7 @@ def test_unpaired_draft_mirror_context_resize( manager.is_draft = True manager.enable_joint_kv_cache_reuse = False manager._mirror_uses_draft_view = uses_draft_view - manager.draft_slot_trail_tokens = trail_tokens + manager._is_dflash_draft = dflash manager.num_extra_kv_tokens = 0 manager.kv_cache_map = {request.py_request_id: kv_cache} manager._resume_and_restore = Mock(return_value=True) From eb3fabcc91f600838a9b906d1be13d6b747b4a61 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Mon, 2 Nov 2026 13:27:53 -0800 Subject: [PATCH 10/14] [TRTLLM-17390][fix] Raise when a DFlash slot seed has no host positions DFlashSpecMetadata.prepare() seeds a fresh generation slot from the host position ids and sequence lengths, which the input preparation always sets; a missing one raises rather than leaving the slot at position 0. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/speculative/dflash.py | 4 +++- .../speculative/hw_agnostic/test_dflash_copy_skip.py | 12 ++++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index afba85439a7e..2edc8068d2d3 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -685,7 +685,9 @@ def prepare(self): ) continue if self.host_position_ids is None or self.seq_lens is None: - continue + raise RuntimeError( + f"DFlash: no host positions to seed the slot of generation request {rid}" + ) if row_starts is None: row_starts = [0] for seq_len in self.seq_lens: diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py index d9b18896ce83..ea6aa4a43045 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py @@ -103,3 +103,15 @@ def counted(dst, src, *args, **kwargs): monkeypatch.undo() assert copies == [] assert worker._batch_to_slot_buffer.tensor[:2].tolist() == [1, 2] + + +def test_a_fresh_generation_slot_without_host_positions_raises() -> None: + # Request 7 is a generation row with no slot yet, so prepare seeds one at + # the row's position; without the host positions it has nothing to seed. + worker = _worker({}) + worker._assign_slot = lambda rid, **kwargs: worker._req_to_slot.setdefault(rid, 3) + metadata = _metadata(worker, [7], num_generations=1) + metadata.host_position_ids = None + metadata.seq_lens = None + with pytest.raises(RuntimeError, match="no host positions"): + _prepare(metadata, [7], num_generations=1) From c1b10c1c4b98287cb8498cd22851cc1e9b38b213 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Mon, 2 Nov 2026 13:28:03 -0800 Subject: [PATCH 11/14] [TRTLLM-17390][fix] Refuse slot-from-0 drafters where a windowed pool is ahead A DFlash worker that numbers slots from 0 (the Kimi K3 workers) refuses to bind a sliding-window draft pool on a disaggregated worker, and raises when it would store a context chunk at slot position 0 for a request past position 0 on such a pool: the manager has released those pages. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/speculative/dflash.py | 31 ++++++++++++ .../kv_cache/test_kvcm2_integration.py | 49 +++++++++++++++++++ .../hw_agnostic/test_dflash_draft_kv.py | 23 +++++++++ 3 files changed, 103 insertions(+) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 2edc8068d2d3..f78bdf48e4a0 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -809,6 +809,9 @@ class DFlashWorker(SpecWorkerBase): # preparation that reads _ctx_block_tables directly sets it False and # numbers every slot from 0. _ctx_reads_rebased_view = True + # Whether the bound draft pool releases pages behind a sliding window. + # Set in _lazy_init_ctx_buffers. + _ctx_pool_windowed = False def __init__( self, @@ -1504,6 +1507,22 @@ def _lazy_init_ctx_buffers( # keeps a page below its windows (kv_cache_manager_v2's # draft_slot_trail_tokens; a V1 pool does not). windows = getattr(draft_kv_cache_manager, "max_attention_window_vec", None) + max_seq_len = getattr(draft_kv_cache_manager, "max_seq_len", None) + self._ctx_pool_windowed = any( + window is not None and (max_seq_len is None or window < max_seq_len) + for window in windows or () + ) + if ( + not self._ctx_reads_rebased_view + and self._ctx_pool_windowed + and getattr(draft_kv_cache_manager, "is_disagg", False) + ): + # A disaggregated generation worker seeds every slot past + # position 0 (_seed_generation_slot). + raise ValueError( + f"{type(self).__name__} numbers drafter slots from 0 and cannot serve a " + "disaggregated worker with a sliding-window draft KV pool." + ) self._ctx_seed_page_aligned = dflash_windows_span_a_page(windows, page_size) and ( all(window is None for window in windows or ()) or getattr(draft_kv_cache_manager, "draft_slot_trail_tokens", 0) >= page_size @@ -1583,6 +1602,18 @@ def _ctx_store_start(self, cur: int, first_pos: int) -> int: and (self._ctx_reuse_addressable or self._ctx_reads_rebased_view) ): return first_pos + if ( + cur == 0 + and first_pos > 0 + and self._ctx_block_tables is not None + and self._ctx_pool_windowed + ): + # The manager has released this request's pages below its window, + # so a slot numbered from 0 would write entries without a page. + raise RuntimeError( + f"{type(self).__name__} numbers drafter slots from 0 and cannot store a context " + f"chunk at position {first_pos} on a sliding-window draft KV pool." + ) return cur def _ctx_base_for_store(self, cur: int, start: int) -> Optional[int]: diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py index ed98c6687995..36c98f5b1d15 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py @@ -2674,6 +2674,55 @@ def own_slots(req: LlmRequest) -> set[int]: draft.shutdown() +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_slot_from_zero_worker_refuses_a_disagg_windowed_draft_pool( + monkeypatch: pytest.MonkeyPatch, is_disagg: bool +) -> None: + """A worker that numbers drafter slots from 0 (the Kimi K3 workers) binds a + sliding-window draft pool, but not on a disaggregated worker, whose every + generation slot starts past position 0.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + monkeypatch.setattr( + "tensorrt_llm._torch.speculative.dflash.validate_dflash_trtllm_gen_runtime", + lambda **kwargs: None, + ) + spec = DFlashDecodingConfig( + max_draft_len=3, speculative_model="draft", attention_backend="TRTLLM", mask_token_id=0 + ) + draft, _, _ = _disagg_draft_fixture(spec, 65536, 512, 256) + draft.is_disagg = is_disagg + try: + worker = DFlashWorker(spec, Mapping()) + worker._ctx_reads_rebased_view = False + draft_model = _PoolBoundDFlashDraftModel(2, 1, 64, 8) + worker.set_draft_model(draft_model) + num_pools = getattr(draft, "num_attention_op_pools", draft.num_pools) + offsets = torch.zeros( + (num_pools, 2, 2, draft.max_blocks_per_seq), dtype=torch.int32, device="cuda" + ) + attn_metadata = SimpleNamespace(max_seq_len=1 << 17, draft_kv_cache_block_offsets=offsets) + + def bind() -> None: + """Bind the drafter's ctx cache to the draft pool.""" + worker._lazy_init_ctx_buffers( + draft_model, + SimpleNamespace(max_num_requests=2), + attn_metadata, + draft_kv_cache_manager=draft, + ) + + if is_disagg: + with pytest.raises(ValueError, match="numbers drafter slots from 0"): + bind() + else: + bind() + assert worker._ctx_block_tables is not None and worker._ctx_pool_windowed + finally: + draft.shutdown() + + @pytest.mark.parametrize("lag", [2, 7, 31, 127]) def test_trailing_drafter_slot_reads_only_pages_its_request_holds( monkeypatch: pytest.MonkeyPatch, lag: int diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py index d89b92753a90..30218d31edd8 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py @@ -86,6 +86,29 @@ def test_dflash_ctx_store_start_follows_the_pool_numbering( assert worker._ctx_store_start(cur, first_pos) == expected +@pytest.mark.parametrize( + ("cur", "first_pos", "raises"), [(0, 8192, True), (0, 0, False), (4096, 4096, False)] +) +def test_slot_from_zero_refuses_a_chunk_past_zero_on_a_windowed_pool( + cur: int, first_pos: int, raises: bool +) -> None: + """A worker that numbers slots from 0 (the Kimi K3 workers) on an unpaired + sliding-window pool cannot store a chunk whose slot starts past 0: the + manager has released the pages that slot would write.""" + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + worker = object.__new__(DFlashWorker) + worker._ctx_block_tables = object() + worker._ctx_reuse_addressable = False + worker._ctx_reads_rebased_view = False + worker._ctx_pool_windowed = True + if raises: + with pytest.raises(RuntimeError, match="numbers drafter slots from 0"): + worker._ctx_store_start(cur, first_pos) + else: + assert worker._ctx_store_start(cur, first_pos) == cur + + @pytest.mark.parametrize( ("cur", "start", "reuse_addressable", "expected"), [ From fa2191bbc92d944c3c98ba1c4f1f31ec0678cd53 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Mon, 2 Nov 2026 13:57:37 -0800 Subject: [PATCH 12/14] [TRTLLM-17390][fix] Keep Kimi K3 fused acceptance off sliding-window draft pools trtllm::k3_spec_accept writes the draft block counts as backed entries, which a sliding-window pool's released entries put below the extent the drafter writes to; the Kimi K3 workers take the torch path on such a pool. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- .../kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/spec_worker.py | 7 +++++-- .../kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py | 7 +++++-- .../modeling_v2/models/kimi_k3_vl/test_k3_decode_worker.py | 1 + 3 files changed, 11 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/spec_worker.py index a35aa480082e..9fe9802ab2c3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/spec_worker.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/spec_worker.py @@ -98,8 +98,8 @@ def _k3_accept_applies( That mode: a decode step (no context requests, at most 8 gen requests) with greedy strict acceptance (no rejection sampling, penalties or guided decoding, the base draft-token and logits layouts), fp32 target - logits of a vocabulary the kernel splits, the V2 Mamba manager's KDA replay record, the draft pool's block - table, and an unsharded bf16 draft embedding. Any other step keeps the torch path. + logits of a vocabulary the kernel splits, the V2 Mamba manager's KDA replay record, the block table of a draft + pool without sliding windows, and an unsharded bf16 draft embedding. Any other step keeps the torch path. """ if not self._k3_accept_step_applies( attn_metadata, spec_metadata, draft_model, num_contexts, num_gens @@ -162,8 +162,11 @@ def _k3_accept_step_applies( or mamba_metadata.state_indices.dtype != torch.int32 ): return False + # The kernel's block counts are backed entries, which a sliding-window + # pool's released entries put below the extent the drafter writes to. if ( self._ctx_block_tables is None + or self._ctx_pool_windowed or getattr(attn_metadata, "draft_kv_cache_block_offsets", None) is None or getattr(attn_metadata, "kv_lens_cuda", None) is None ): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py index a35aa480082e..9fe9802ab2c3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py @@ -98,8 +98,8 @@ def _k3_accept_applies( That mode: a decode step (no context requests, at most 8 gen requests) with greedy strict acceptance (no rejection sampling, penalties or guided decoding, the base draft-token and logits layouts), fp32 target - logits of a vocabulary the kernel splits, the V2 Mamba manager's KDA replay record, the draft pool's block - table, and an unsharded bf16 draft embedding. Any other step keeps the torch path. + logits of a vocabulary the kernel splits, the V2 Mamba manager's KDA replay record, the block table of a draft + pool without sliding windows, and an unsharded bf16 draft embedding. Any other step keeps the torch path. """ if not self._k3_accept_step_applies( attn_metadata, spec_metadata, draft_model, num_contexts, num_gens @@ -162,8 +162,11 @@ def _k3_accept_step_applies( or mamba_metadata.state_indices.dtype != torch.int32 ): return False + # The kernel's block counts are backed entries, which a sliding-window + # pool's released entries put below the extent the drafter writes to. if ( self._ctx_block_tables is None + or self._ctx_pool_windowed or getattr(attn_metadata, "draft_kv_cache_block_offsets", None) is None or getattr(attn_metadata, "kv_lens_cuda", None) is None ): diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_decode_worker.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_decode_worker.py index 440918ec1410..55e00273903f 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_decode_worker.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_decode_worker.py @@ -239,6 +239,7 @@ def mutate(case): "attn.mamba_metadata.state_indices", torch.zeros(4, dtype=torch.int64) ), "the private context arena": _set("worker._ctx_block_tables", None), + "a sliding-window draft pool": _set("worker._ctx_pool_windowed", True), "no draft block offsets": _set("attn.draft_kv_cache_block_offsets", None), "no KV lengths": _set("attn.kv_lens_cuda", None), "a vocab-sharded draft embedding": _set( From 4d8ae992d7fa19d1182338e805cda12814642b21 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Mon, 2 Nov 2026 14:08:46 -0800 Subject: [PATCH 13/14] [TRTLLM-17390][fix] Let a slot-from-0 drafter restart a chunk within the trail The context guard for a drafter that numbers slots from 0 raises only when the chunk starts past the pool's trailing page, where the slot's window would outgrow the pages the manager keeps; within it every entry stays backed. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/speculative/dflash.py | 15 ++++++++++----- .../hw_agnostic/test_dflash_draft_kv.py | 9 ++++++--- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index f78bdf48e4a0..75a76192347e 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -809,9 +809,11 @@ class DFlashWorker(SpecWorkerBase): # preparation that reads _ctx_block_tables directly sets it False and # numbers every slot from 0. _ctx_reads_rebased_view = True - # Whether the bound draft pool releases pages behind a sliding window. - # Set in _lazy_init_ctx_buffers. + # Whether the bound draft pool releases pages behind a sliding window, and + # how many tokens it keeps below each window (kv_cache_manager_v2's + # draft_slot_trail_tokens). Set in _lazy_init_ctx_buffers. _ctx_pool_windowed = False + _ctx_pool_trail = 0 def __init__( self, @@ -1512,6 +1514,7 @@ def _lazy_init_ctx_buffers( window is not None and (max_seq_len is None or window < max_seq_len) for window in windows or () ) + self._ctx_pool_trail = getattr(draft_kv_cache_manager, "draft_slot_trail_tokens", 0) if ( not self._ctx_reads_rebased_view and self._ctx_pool_windowed @@ -1604,12 +1607,14 @@ def _ctx_store_start(self, cur: int, first_pos: int) -> int: return first_pos if ( cur == 0 - and first_pos > 0 + and first_pos > self._ctx_pool_trail and self._ctx_block_tables is not None and self._ctx_pool_windowed ): - # The manager has released this request's pages below its window, - # so a slot numbered from 0 would write entries without a page. + # A slot numbered from 0 lags the request's absolute positions by + # first_pos, and the manager keeps pages only from the absolute + # position less the window and trail: past the trail, the slot's + # window reaches entries without a page once the request outgrows it. raise RuntimeError( f"{type(self).__name__} numbers drafter slots from 0 and cannot store a context " f"chunk at position {first_pos} on a sliding-window draft KV pool." diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py index 30218d31edd8..e1eb02dc5809 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py @@ -87,14 +87,16 @@ def test_dflash_ctx_store_start_follows_the_pool_numbering( @pytest.mark.parametrize( - ("cur", "first_pos", "raises"), [(0, 8192, True), (0, 0, False), (4096, 4096, False)] + ("cur", "first_pos", "raises"), + [(0, 8192, True), (0, 65, True), (0, 64, False), (0, 0, False), (4096, 4096, False)], ) def test_slot_from_zero_refuses_a_chunk_past_zero_on_a_windowed_pool( cur: int, first_pos: int, raises: bool ) -> None: """A worker that numbers slots from 0 (the Kimi K3 workers) on an unpaired - sliding-window pool cannot store a chunk whose slot starts past 0: the - manager has released the pages that slot would write.""" + sliding-window pool with a 64-token trail cannot store a chunk whose slot + starts more than the trail past 0: the manager releases the pages that slot + would address. Within the trail every entry stays on a page it holds.""" from tensorrt_llm._torch.speculative.dflash import DFlashWorker worker = object.__new__(DFlashWorker) @@ -102,6 +104,7 @@ def test_slot_from_zero_refuses_a_chunk_past_zero_on_a_windowed_pool( worker._ctx_reuse_addressable = False worker._ctx_reads_rebased_view = False worker._ctx_pool_windowed = True + worker._ctx_pool_trail = 64 if raises: with pytest.raises(RuntimeError, match="numbers drafter slots from 0"): worker._ctx_store_start(cur, first_pos) From 216eddbd2e3b4485b88a5e8352d9a7c7c7695a68 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Mon, 2 Nov 2026 14:41:21 -0800 Subject: [PATCH 14/14] [TRTLLM-17390][fix] Skip a slot-from-0 drafter's context store past the trail A drafter that numbers slots from 0 cannot store a context chunk more than a sliding-window pool's trailing page past position 0. That request's context store is skipped and its slot released with a warning, like a context overflow, instead of failing the batch's forward. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/speculative/dflash.py | 27 ++++-- .../hw_agnostic/test_dflash_draft_kv.py | 84 ++++++++++++++++--- 2 files changed, 92 insertions(+), 19 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 75a76192347e..e9e3e5e8a66f 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -1582,8 +1582,9 @@ def _ctx_paged_index_args(self): indptr = self._ctx_kv_indptr return table.flatten().contiguous(), indptr - def _ctx_store_start(self, cur: int, first_pos: int) -> int: - """Position at which a context chunk's drafter K/V is stored. + def _ctx_store_start(self, cur: int, first_pos: int) -> Optional[int]: + """Position at which a context chunk's drafter K/V is stored, or None + when this worker cannot store it. ``cur`` is the slot's stored length, ``first_pos`` the chunk's first target position. They differ only when the target skipped @@ -1596,7 +1597,8 @@ def _ctx_store_start(self, cur: int, first_pos: int) -> int: an unpaired pool when the worker reads the rebased view, which hides the empty prefix. Otherwise (the private arena, or an unpaired pool read from position 0) every request numbers from 0, where the tail is - all the context there is. + all the context there is -- except on a sliding-window pool more than + its trailing page past position 0 (None). """ if ( cur == 0 @@ -1615,10 +1617,7 @@ def _ctx_store_start(self, cur: int, first_pos: int) -> int: # first_pos, and the manager keeps pages only from the absolute # position less the window and trail: past the trail, the slot's # window reaches entries without a page once the request outgrows it. - raise RuntimeError( - f"{type(self).__name__} numbers drafter slots from 0 and cannot store a context " - f"chunk at position {first_pos} on a sliding-window draft KV pool." - ) + return None return cur def _ctx_base_for_store(self, cur: int, start: int) -> Optional[int]: @@ -1953,6 +1952,20 @@ def _store_prefill_context( # (precompute_context_kv is per-token in (hidden, p)); an unpaired # one does not, so the slot reads from a base above it. start = self._ctx_store_start(cur, first_pos) + if start is None: + # Request-level, like the overflow below: the slot is released + # rather than failing the batch's forward. + logger.warning( + f"DFlash: {type(self).__name__} numbers drafter slots from 0 and cannot store " + f"the context chunk at position {first_pos} of request {req_id} on a " + "sliding-window draft KV pool; skipping its context store." + ) + self._req_to_slot.pop(req_id, None) + self._req_ctx_pos.pop(req_id, None) + ctx_len_updates[slot] = 0 + self._free_slots.append(slot) + offset += slen + continue base = self._ctx_base_for_store(cur, start) if base is not None: ctx_base_updates[slot] = base diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py index e1eb02dc5809..57635baaa5ab 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py @@ -6,6 +6,9 @@ buffers and must not allocate a redundant manager or be gated on its capacity. """ +from types import SimpleNamespace +from typing import Optional + import pytest from tensorrt_llm._torch.speculative.interface import should_use_separate_draft_kv_cache @@ -87,29 +90,86 @@ def test_dflash_ctx_store_start_follows_the_pool_numbering( @pytest.mark.parametrize( - ("cur", "first_pos", "raises"), - [(0, 8192, True), (0, 65, True), (0, 64, False), (0, 0, False), (4096, 4096, False)], + ("cur", "first_pos", "expected"), + [(0, 8192, None), (0, 65, None), (0, 64, 0), (0, 0, 0), (4096, 4096, 4096)], ) -def test_slot_from_zero_refuses_a_chunk_past_zero_on_a_windowed_pool( - cur: int, first_pos: int, raises: bool +def test_slot_from_zero_cannot_store_a_chunk_past_the_trail_on_a_windowed_pool( + cur: int, first_pos: int, expected: Optional[int] ) -> None: """A worker that numbers slots from 0 (the Kimi K3 workers) on an unpaired - sliding-window pool with a 64-token trail cannot store a chunk whose slot - starts more than the trail past 0: the manager releases the pages that slot - would address. Within the trail every entry stays on a page it holds.""" + sliding-window pool with a 64-token trail has no store start for a chunk + whose slot starts more than the trail past 0: the manager releases the + pages that slot would address. Within the trail every entry stays on a page + it holds.""" + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + worker = object.__new__(DFlashWorker) + worker._ctx_block_tables = object() + worker._ctx_reuse_addressable = False + worker._ctx_reads_rebased_view = False + worker._ctx_pool_windowed = True + worker._ctx_pool_trail = 64 + assert worker._ctx_store_start(cur, first_pos) == expected + + +def test_slot_from_zero_skips_a_context_store_past_the_trail( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A Kimi K3 style worker whose request lost its slot mid-prompt skips that + chunk's context store and releases the slot, with a warning, instead of + failing the batch.""" + import collections + + import torch + + from tensorrt_llm._torch.speculative import dflash from tensorrt_llm._torch.speculative.dflash import DFlashWorker worker = object.__new__(DFlashWorker) worker._ctx_block_tables = object() + worker._ctx_block_counts = torch.tensor([300]) + worker._ctx_page_size = 64 + worker._max_ctx = 1 << 17 worker._ctx_reuse_addressable = False worker._ctx_reads_rebased_view = False worker._ctx_pool_windowed = True worker._ctx_pool_trail = 64 - if raises: - with pytest.raises(RuntimeError, match="numbers drafter slots from 0"): - worker._ctx_store_start(cur, first_pos) - else: - assert worker._ctx_store_start(cur, first_pos) == cur + worker._req_ctx_pos = {} + worker._req_to_slot = {} + worker._free_slots = collections.deque([2]) + worker._ctx_len_host = [0, 0, 0] + worker._ctx_base_host = [0, 0, 0] + worker._ctx_len = torch.zeros(3, dtype=torch.long) + worker._ctx_base_pages = torch.zeros(3, dtype=torch.long) + warnings = [] + monkeypatch.setattr(dflash.logger, "warning", warnings.append) + + slen, first_pos = 4, 8192 + draft_model = SimpleNamespace( + fc=object(), hidden_norm=object(), project_target_hidden=lambda hs: hs + ) + spec_metadata = SimpleNamespace( + get_hidden_states=lambda n: torch.zeros(n, 8), + host_position_ids=torch.arange(first_pos, first_pos + slen), + request_ids=[7], + ) + attn_metadata = SimpleNamespace( + num_ctx_tokens=slen, num_contexts=1, _seq_lens=torch.tensor([slen]) + ) + worker._store_prefill_context( + draft_model, + spec_metadata, + attn_metadata, + torch.arange(first_pos, first_pos + slen), + slen, + ) + + assert 7 not in worker._req_to_slot and 7 not in worker._req_ctx_pos + assert list(worker._free_slots) == [2] + assert worker._ctx_len_host[2] == 0 and worker._ctx_len.tolist()[2] == 0 + assert len(warnings) == 1 + assert "skipping its context store" in warnings[0] + assert "drafts nothing" not in warnings[0] @pytest.mark.parametrize(