diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py index c7ca1891b642..b46d8478ffd6 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py @@ -1301,8 +1301,13 @@ def copy_batch_block_offsets( beam_width: int, num_contexts: int, num_seqs: int, + max_blocks: Optional[int] = None, ) -> None: - """For compatibility with AttentionOp, copy only the SWA block offsets.""" + """For compatibility with AttentionOp, copy only the SWA block offsets. + + max_blocks is accepted for signature parity with KVCacheManager; the + copy below is already bounded by the precomputed SWA table width. + """ assert beam_width == 1, "DSV4 only supports beam width 1 now" assert dst_tensor.is_cuda, "copy_batch_block_offsets expects a CUDA destination" dst_tensor.fill_(BAD_PAGE_INDEX) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 931830253cba..40d15970398b 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -31,6 +31,7 @@ from tensorrt_llm._utils import get_sm_version, maybe_pin_memory, prefer_pinned from tensorrt_llm.bindings.internal import thop from tensorrt_llm.functional import AttentionMaskType +from tensorrt_llm.math_utils import ceil_div from tensorrt_llm.models.modeling_utils import QuantConfig from ..utils import (compute_swizzled_sf_shape, get_global_attrs, @@ -586,24 +587,47 @@ def prepare(self) -> None: # kv block offsets assert self.request_ids is not None if self.kv_cache_manager is not None: - self.kv_cache_manager.copy_batch_block_offsets( - self.kv_cache_block_offsets, self.request_ids, self.beam_width, - self.num_contexts, self.num_seqs) - - error_message = ( - f"The max KV cache length of input sequences ({self.kv_lens[:self.num_seqs].max()}) " + max_kv_len = int(self.kv_lens[:self.num_seqs].max()) + assert max_kv_len <= self.kv_cache_manager.max_seq_len, ( + f"The max KV cache length of input sequences ({max_kv_len}) " f"exceeds the KV cache manager's maximum supported length " f"({self.kv_cache_manager.max_seq_len}).") - assert self.kv_lens[:self.num_seqs].max( - ) <= self.kv_cache_manager.max_seq_len, error_message + # On the non-speculative path the host kv_lens snapshot bounds + # every block-table access, so the staged/H2D width can be capped + # at the batch's maximum instead of max_seq_len's worth of + # columns. Speculative decoding must stage the full width: + # draft/tree sub-steps and the overlap scheduler advance + # kv_lens_cuda on device past the host snapshot, and their + # kernels dereference block columns a host-derived cap would + # leave unstaged (uninitialized in this buffer). + spec_active = (self.draft_kv_cache_manager is not None + or self.is_spec_decoding_enabled + or bool(self.kv_cache_params.num_extra_kv_tokens) or + (self.runtime_features is not None and + self.runtime_features.has_speculative_draft_tokens)) + max_blocks = None + if not spec_active and self.kv_cache_manager.tokens_per_block: + max_blocks = ceil_div(max_kv_len, + self.kv_cache_manager.tokens_per_block) + self.kv_cache_manager.copy_batch_block_offsets( + self.kv_cache_block_offsets, + self.request_ids, + self.beam_width, + self.num_contexts, + self.num_seqs, + max_blocks=max_blocks) # Also prepare draft KV cache block offsets if draft_kv_cache_manager exists if self.draft_kv_cache_manager is not None: # Use the wrapper method which works for both V1 and V2 self.draft_kv_cache_manager.copy_batch_block_offsets( - self.draft_kv_cache_block_offsets, self.request_ids, - self.beam_width, self.num_contexts, self.num_seqs) + self.draft_kv_cache_block_offsets, + self.request_ids, + self.beam_width, + self.num_contexts, + self.num_seqs, + max_blocks=max_blocks) # Don't pass self.kv_lens as kv_lens here because it includes extra # tokens. Use the actual KV length (without extra tokens) for diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 5cab7283f033..9c323799a77a 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -369,9 +369,14 @@ def prepare(self, attn_metadata: AttentionMetadata): self.state_indices_cpu[:batch_size], non_blocking=True) else: # indices is a Python sequence (e.g. List[int]); data - # already lives on host, CPU staging is fine. - for i, idx in enumerate(indices): - self.state_indices_cpu[i] = idx + # already lives on host, CPU staging is fine. One bulk + # conversion instead of a per-element tensor write. + assert len(indices) == batch_size, ( + f"get_state_indices() returned {len(indices)} entries for " + f"a batch of {batch_size} requests.") + self.state_indices_cpu[:batch_size].copy_( + torch.as_tensor(indices, + dtype=self.state_indices_cpu.dtype)) self.state_indices[:batch_size].copy_( self.state_indices_cpu[:batch_size], non_blocking=True) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 46d3c55a0771..ffe7b6f5f83e 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -3194,7 +3194,10 @@ def copy_batch_block_offsets( beam_width: int, num_contexts: int, num_seqs: int, + max_blocks: Optional[int] = None, ): + # max_blocks is accepted for signature parity with KVCacheManager; the + # device-side copy op here already scales with allocated blocks only. assert beam_width == 1, "beam_width must be 1 for KVCacheManagerV2" copy_idx = self.index_mapper.get_copy_index(request_ids, num_contexts, beam_width) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index df137d96de6d..d02e19d5b95c 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -1984,15 +1984,18 @@ def _setup_state_indices(self, requests=None) -> None: self.cuda_state_indices.copy_(self._host_state_indices, non_blocking=True) + is_dummy = [req.is_dummy for req in requests] self._refresh_dummy_request_mask( - [req.is_dummy for req in self.requests]) + is_dummy if requests is + self.requests else [req.is_dummy for req in self.requests]) # Build request_id → pool block offset mapping so that # get_state_indices can return indices in arbitrary request order. - for i, req in enumerate(requests): - self._request_id_to_state_index[ - req.py_request_id] = self._host_state_indices[i].item() - self._request_id_to_is_dummy[req.py_request_id] = req.is_dummy + # Bulk tolist avoids a per-request tensor-index + .item() round-trip. + state_values = self._host_state_indices[:n].tolist() + for req, value, dummy in zip(requests, state_values, is_dummy): + self._request_id_to_state_index[req.py_request_id] = value + self._request_id_to_is_dummy[req.py_request_id] = dummy def get_state_indices(self, request_ids: Optional[List[int]] = None, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 62a442972830..dff5b770704a 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -675,6 +675,18 @@ def __init__( self.position_ids_cuda = torch.empty((self.max_num_tokens, ), dtype=torch.int, device='cuda') + # Steady-state generation-only prepare cache (non-speculative overlap + # decode). Holds the per-request lists that are invariant while the + # scheduled generation batch keeps the same composition, plus a pinned + # cached-token counter advanced by one per step (host-side bookkeeping + # only; the device position buffer is advanced in place and this + # buffer is never the source of an async H2D). Invalidated (set to + # None) by every full _prepare_tp_inputs pass. + self._steady_gen_cache: Optional[Dict[str, Any]] = None + self._steady_gen_positions_pinned = torch.empty( + (self.max_num_tokens, ), + dtype=torch.int, + pin_memory=prefer_pinned()) if self.use_mrope: self.mrope_position_ids_cuda = torch.empty( (3, 1, self.max_num_tokens), dtype=torch.int, device='cuda') @@ -3269,6 +3281,143 @@ def _apply_incremental_update_target( return inputs, self.gather_ids_cuda[:num_generation_tokens] + def _can_use_steady_gen_fast_prepare( + self, scheduled_requests: ScheduledRequests, + new_tokens_device: Optional[torch.Tensor], + next_draft_tokens_device: Optional[torch.Tensor], + spec_metadata: Optional[SpecMetadata]) -> bool: + """Check whether the cached steady-state generation prepare applies. + + The cache is only recorded by a full _prepare_tp_inputs pass whose + batch consisted purely of non-dummy generation requests that all had + a previous overlap-scheduler tensor (see the recording site), so the + per-step check only needs to confirm the dynamic conditions: still a + generation-only batch with the exact same requests in the same order. + """ + cache = self._steady_gen_cache + if cache is None or self.is_warmup: + return False + if new_tokens_device is None or next_draft_tokens_device is not None \ + or spec_metadata is not None: + return False + if scheduled_requests.num_context_requests > 0: + return False + generation_requests = scheduled_requests.generation_requests + if len(generation_requests) != cache['num_requests']: + return False + return cache['request_ids'] == [ + request.py_request_id for request in generation_requests + ] + + @nvtx_range("_apply_steady_gen_fast_prepare") + def _apply_steady_gen_fast_prepare( + self, kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2], + attn_metadata: AttentionMetadata, + new_tensors_device: SampleStateTensors, + resource_manager: Optional[ResourceManager]): + """Prepare inputs for an unchanged generation-only batch. + + Every request advanced by exactly one committed token since the last + prepare, so instead of re-walking the batch in Python this advances + the cached positions in place (device position buffer plus a pinned + host counter), reuses the seq-slot buffer already on device, and + refreshes only the per-step metadata. For mrope models (recorded only + for batches with no actual mrope work) the (3,1,N) broadcast buffer + the model reads is the one advanced. + """ + cache = self._steady_gen_cache + num_requests = cache['num_requests'] + + # Positions and cached-token counts are the same values in this + # regime; advance both by one. The device-side position buffer is + # advanced in place: it still holds the previous step's positions + # because only _prepare_tp_inputs writes it and the cache validity + # invariant guarantees the previous pass wrote these same rows. This + # avoids reusing a mutated pinned buffer as the source of an async + # H2D whose previous-step copy may still be pending under the overlap + # scheduler (the nvbug 6293536 hazard class; see + # KVCacheManager._stage_block_offsets_for_copy). The pinned buffer is + # host-side bookkeeping only. + use_mrope = cache['use_mrope'] + positions = self._steady_gen_positions_pinned[:num_requests] + positions.add_(1) + if use_mrope: + # Text-only batch on an mrope model: the recording pass broadcast + # the scalar positions onto all three axes of the (3,1,N) buffer, + # which is what the model (and any captured CUDA graph) reads, so + # advance it in place. position_ids_cuda is reseeded by the next + # full pass. + self.mrope_position_ids_cuda[:, :, :num_requests].add_(1) + else: + self.position_ids_cuda[:num_requests].add_(1) + num_cached_tokens_per_seq = positions.tolist() + + # Gather this step's input tokens from the previous iteration's device + # sample buffer; the seq-slot indices in previous_batch_indices_cuda + # are unchanged since the last full pass. + previous_slots = self.previous_batch_indices_cuda[:num_requests] + new_tokens = new_tensors_device.new_tokens[:1, previous_slots, :self. + max_beam_width] + self.input_ids_cuda[:num_requests * self.max_beam_width].copy_( + new_tokens.flatten(), non_blocking=True) + + if not attn_metadata.is_cuda_graph: + attn_metadata.seq_lens = cache['seq_lens_ones'] + attn_metadata.beam_width = 1 + attn_metadata.request_ids = cache['request_ids'] + attn_metadata.prompt_lens = cache['prompt_lens'] + attn_metadata.num_contexts = 0 + attn_metadata.num_chunked_ctx_requests = 0 + attn_metadata.kv_cache_params = KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=num_cached_tokens_per_seq, + num_extra_kv_tokens=get_num_extra_kv_tokens(None)) + attn_metadata.kv_cache_manager = kv_cache_manager + if hasattr(self.model.model_config.pretrained_config, 'chunk_size'): + attn_metadata.mamba_chunk_size = \ + self.model.model_config.pretrained_config.chunk_size + with nvtx_range("steady_gen_metadata_prepare"): + attn_metadata.prepare() + + attn_all_rank_num_tokens = self._get_all_rank_num_tokens(attn_metadata) + padded_num_tokens, can_run_piecewise_cuda_graph, attn_all_rank_num_tokens = \ + self._get_padding_params(num_requests, 0, attn_all_rank_num_tokens) + set_per_request_piecewise_cuda_graph_flag(can_run_piecewise_cuda_graph) + attn_metadata.padded_num_tokens = ( + padded_num_tokens if padded_num_tokens != num_requests else None) + virtual_num_tokens = num_requests + if attn_metadata.padded_num_tokens is not None: + self.input_ids_cuda[num_requests:padded_num_tokens].fill_(0) + # Zero-fill the padding tail of whichever position layout the + # model consumes, matching the full pass. + if use_mrope: + self.mrope_position_ids_cuda[:, :, num_requests: + padded_num_tokens].fill_(0) + else: + self.position_ids_cuda[num_requests:padded_num_tokens].fill_(0) + virtual_num_tokens = padded_num_tokens + + self.iter_states['num_ctx_requests'] = 0 + self.iter_states['num_ctx_tokens'] = 0 + self.iter_states['num_generation_tokens'] = num_requests + self.iter_states['cached_kv_tokens'] = sum(num_cached_tokens_per_seq) + + if use_mrope: + final_position_ids = \ + self.mrope_position_ids_cuda[:, :, :virtual_num_tokens] + else: + final_position_ids = \ + self.position_ids_cuda[:virtual_num_tokens].unsqueeze(0) + inputs = { + 'attn_metadata': attn_metadata, + 'input_ids': self.input_ids_cuda[:virtual_num_tokens], + 'position_ids': final_position_ids, + 'inputs_embeds': None, + 'multimodal_params': [], + 'resource_manager': resource_manager, + } + return inputs, None + def _prepare_tp_inputs( self, scheduled_requests: ScheduledRequests, @@ -3305,12 +3454,28 @@ def _prepare_tp_inputs( if self._can_use_incremental_update(scheduled_requests, new_tokens_device, next_draft_tokens_device): + # Spec engines never record the steady-gen cache, but invalidate + # defensively so the two fast paths can never interleave if the + # gates ever evolve. + self._steady_gen_cache = None return self._apply_incremental_update( scheduled_requests, kv_cache_manager, attn_metadata, spec_metadata, new_tensors_device, cache_indirection_buffer, num_accepted_tokens_device, req_id_to_old_request, resource_manager) + if self._can_use_steady_gen_fast_prepare(scheduled_requests, + new_tokens_device, + next_draft_tokens_device, + spec_metadata): + return self._apply_steady_gen_fast_prepare(kv_cache_manager, + attn_metadata, + new_tensors_device, + resource_manager) + # Any full pass invalidates the steady-state cache; it is re-recorded + # at the end of this pass when the batch qualifies. + self._steady_gen_cache = None + # Hoist self.use_mrope to a function-scope local so the per-request / # per-context-request mrope branches use LOAD_FAST instead of LOAD_ATTR. _use_mrope = self.use_mrope @@ -4226,6 +4391,12 @@ def previous_seq_slots_device(): if hasattr(self.model.model_config.pretrained_config, 'chunk_size'): attn_metadata.mamba_chunk_size = self.model.model_config.pretrained_config.chunk_size + # Some sparse backends (RocketKV) clamp + # kv_cache_params.num_cached_tokens_per_seq in place during prepare(), + # and KVCacheParams holds the list by reference. Snapshot the true + # pre-prepare counts so the steady-gen recording below stores values + # that the per-step prepare() can re-clamp from scratch. + num_cached_tokens_snapshot = list(num_cached_tokens_per_seq) attn_metadata.prepare() cross_attention_inputs = (self._prepare_enc_dec_cross_attn_inputs( cross_encoder_hidden_states, @@ -4354,6 +4525,50 @@ def previous_seq_slots_device(): self.previous_request_ids = all_gen_request_ids self.has_previous_device_draft = next_draft_tokens_device is not None + # Record the steady-state generation cache when this pass handled + # purely non-dummy generation requests that all carried a previous + # overlap-scheduler tensor (previous_batch_len == _n_gen implies + # every request took that branch and none appended input_ids). + # While the batch composition holds, the next passes only need to + # advance positions by one and refresh per-step metadata. + # MRoPE models are supported only for batches with no actual mrope + # work (text-only requests, empty mrope lists below): the full + # pass routes use_mrope models through the (3,1,N) + # mrope_position_ids_cuda layout even then (to keep torch.compile + # guards stable), with all three axes equal to the scalar + # positions, so the fast path advances that buffer in place and + # returns the same layout (see _apply_steady_gen_fast_prepare). + if (self.spec_config is None and not self.is_draft_model + and spec_metadata is None and new_tokens_device is not None + and self.guided_decoder is None + and not self.enable_attention_dp and not mrope_position_ids + and not mrope_delta_write_seq_slots + and not mrope_delta_read_seq_slots + and not self.use_beam_search and self.max_beam_width == 1 + and not is_enc_dec and not _has_cp_helix + and num_ctx_requests == 0 and not extend_requests + and not first_draft_requests and _n_gen > 0 + and previous_batch_len == _n_gen and num_tokens == 0 + and not _has_any_multimodal_request + and not multimodal_params_list and not lora_params + and attn_metadata.padded_num_tokens is None + and self._get_position_id_offset() == 0): + self._steady_gen_positions_pinned[:_n_gen].copy_( + torch.as_tensor(num_cached_tokens_snapshot, + dtype=torch.int)) + self._steady_gen_cache = { + 'num_requests': + _n_gen, + 'request_ids': + all_gen_request_ids, + 'prompt_lens': + prompt_lengths, + 'seq_lens_ones': + maybe_pin_memory(torch.ones(_n_gen, dtype=torch.int)), + 'use_mrope': + _use_mrope, + } + return inputs, self.gather_ids_cuda[:len( gather_ids)] if self.enable_spec_decode else None diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 5bac35ab06d0..9bcbf9cd57bb 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -2173,9 +2173,13 @@ def _validate_and_adjust_attention_windows( def pin_blocks(self, request_id: int): self.impl.pin_blocks(request_id) - def copy_batch_block_offsets(self, dst_tensor: torch.Tensor, - request_ids: List[int], beam_width: int, - num_context: int, num_seqs: int): + def copy_batch_block_offsets(self, + dst_tensor: torch.Tensor, + request_ids: List[int], + beam_width: int, + num_context: int, + num_seqs: int, + max_blocks: Optional[int] = None): # Fill the persistent host buffer in place, exactly as before. CPU-side # consumers read self.host_kv_cache_block_offsets directly and depend on # its persistent, max_batch-sized layout: DSA sparse attention, the @@ -2241,23 +2245,39 @@ def copy_batch_block_offsets(self, dst_tensor: torch.Tensor, # matching the already-safe kv_lens / block_ids_per_seq staging. The # persistent buffer above is untouched by this and stays valid for the # synchronous CPU readers. - host_block_offsets = self._stage_block_offsets_for_copy(num_seqs) + host_block_offsets = self._stage_block_offsets_for_copy( + num_seqs, max_blocks) + width = host_block_offsets.shape[-1] for pool_idx in range(self.num_pools): - dst_tensor[pool_idx, :num_seqs].copy_(host_block_offsets[pool_idx], - non_blocking=True) + dst_tensor[pool_idx, :num_seqs, :, :width].copy_( + host_block_offsets[pool_idx], non_blocking=True) - def _stage_block_offsets_for_copy(self, num_rows: int) -> torch.Tensor: + def _stage_block_offsets_for_copy( + self, + num_rows: int, + max_blocks: Optional[int] = None) -> torch.Tensor: """Snapshot the first ``num_rows`` rows of the persistent host block offset buffer into a fresh pinned buffer, to serve as the private source - of an asynchronous H2D copy (nvbug 6293536).""" + of an asynchronous H2D copy (nvbug 6293536). + + ``max_blocks`` bounds the copied block width. The buffer is laid out + for max_seq_len (max_blocks_per_seq columns) but consumers only read + each sequence's allocated block prefix, so a caller that knows the + batch's maximum KV length can skip the unused tail — with a large + max_seq_len the tail dominates the copy cost.""" + if max_blocks is None: + width = self.max_blocks_per_seq + else: + width = min(max(max_blocks, 1), self.max_blocks_per_seq) host_block_offsets = torch.empty(self.num_pools, num_rows, 2, - self.max_blocks_per_seq, + width, dtype=torch.int32, pin_memory=prefer_pinned(), device='cpu') - host_block_offsets.copy_(self.host_kv_cache_block_offsets[:, :num_rows]) + host_block_offsets.copy_( + self.host_kv_cache_block_offsets[:, :num_rows, :, :width]) return host_block_offsets def truncate_blocks(self, target_tokens: List[int], diff --git a/tests/unittest/_torch/executor/test_kv_block_offset_overlap_race.py b/tests/unittest/_torch/executor/test_kv_block_offset_overlap_race.py index 92aba91ead1b..ff5743430938 100644 --- a/tests/unittest/_torch/executor/test_kv_block_offset_overlap_race.py +++ b/tests/unittest/_torch/executor/test_kv_block_offset_overlap_race.py @@ -102,6 +102,53 @@ def _reference_offsets(mgr, ids): return dst[:, : len(ids)].clone() +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device") +def test_copy_batch_block_offsets_max_blocks_staging_width(): + """``max_blocks`` bounds the staged H2D width; ``None`` stages the full + allocated width. + + The bounded copy must leave destination columns beyond the cap untouched + (callers only pass a cap when nothing reads past it), and the unbounded + copy must ship every allocated block: speculative decoding allocates + blocks past the host kv_lens snapshot and its kernels dereference those + columns, so capping by a host-derived block count corrupted the device + block table (EAGLE3 warmup illegal memory access). + """ + mgr = _build_manager() + + ids = list(range(1, 1 + _BATCH)) + toks = [_TOKENS_PER_SEQ] * _BATCH # 5 allocated blocks per sequence + mgr.add_dummy_requests(request_ids=ids, token_nums=toks, prepare_resource=True) + allocated_blocks = _TOKENS_PER_SEQ // _TOKENS_PER_BLOCK + + ref = _reference_offsets(mgr, ids) + + sentinel = -12345 + capped_width = 2 + assert capped_width < allocated_blocks <= mgr.max_blocks_per_seq + + dst = torch.full( + (mgr.num_pools, 2 * _BATCH, 2, mgr.max_blocks_per_seq), + sentinel, + dtype=torch.int32, + device="cuda", + ) + mgr.copy_batch_block_offsets(dst, ids, 1, _BATCH, _BATCH, max_blocks=capped_width) + torch.cuda.synchronize() + assert torch.equal(dst[:, :_BATCH, :, :capped_width], ref[..., :capped_width]) + assert (dst[:, :_BATCH, :, capped_width:] == sentinel).all(), ( + "bounded staging must not write past the requested block width" + ) + + dst.fill_(sentinel) + mgr.copy_batch_block_offsets(dst, ids, 1, _BATCH, _BATCH, max_blocks=None) + torch.cuda.synchronize() + assert torch.equal(dst[:, :_BATCH, :, :allocated_blocks], ref[..., :allocated_blocks]), ( + "max_blocks=None must stage every allocated block, including blocks " + "past the batch's current kv length (speculative decoding reads them)" + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device") def test_copy_batch_block_offsets_survives_overlap_overwrite(): mgr = _build_manager() diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 6298565205eb..4e7faecbbea2 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -207,6 +207,81 @@ def test_kv_lens_runtime_with_eagle3_one_model(): f"kv_lens should be {expected_kv_lens_with_extra.tolist()}, but got {kv_lens_internal.tolist()}" +def _make_mock_kv_cache_manager(num_seqs: int) -> MagicMock: + mock_kv_cache_manager = MagicMock() + mock_kv_cache_manager.tokens_per_block = 32 + mock_kv_cache_manager.num_pools = 1 + mock_kv_cache_manager.num_attention_op_pools = 1 + mock_kv_cache_manager.max_blocks_per_seq = 16 + mock_kv_cache_manager.max_batch_size = num_seqs + mock_kv_cache_manager.max_seq_len = 512 + mock_kv_cache_manager.copy_batch_block_offsets = MagicMock() + return mock_kv_cache_manager + + +@pytest.mark.parametrize("spec_signal", [ + None, "num_extra_kv_tokens", "is_spec_decoding_enabled", + "has_speculative_draft_tokens", "draft_kv_cache_manager" +]) +def test_block_offsets_staging_width_spec_gate(spec_signal): + """prepare() caps the staged block-table width by the batch's max KV + length only on the non-speculative path. + + Any speculative-decoding signal must disable the cap (max_blocks=None): + spec kernels address block columns past the host kv_lens snapshot + (device-side kv_lens advances in draft/tree sub-steps, draft-token blocks + are allocated ahead), so a host-derived cap leaves columns they + dereference unstaged. Regression test for the EAGLE3 warmup illegal + memory access. + """ + num_seqs = 3 + prompt_lens = [50, 100, 75] + seq_lens_q = [1, 1, 1] + num_cached_tokens_per_seq = [ + prompt_lens[i] - seq_lens_q[i] for i in range(num_seqs) + ] + + mock_kv_cache_manager = _make_mock_kv_cache_manager(num_seqs) + metadata_kwargs = dict( + max_num_requests=num_seqs, + max_num_tokens=sum(seq_lens_q), + kv_cache_manager=mock_kv_cache_manager, + ) + mock_draft_manager = None + if spec_signal == "draft_kv_cache_manager": + mock_draft_manager = _make_mock_kv_cache_manager(num_seqs) + metadata_kwargs["draft_kv_cache_manager"] = mock_draft_manager + + attn_metadata = TrtllmAttentionMetadata(**metadata_kwargs) + if spec_signal == "is_spec_decoding_enabled": + attn_metadata.is_spec_decoding_enabled = True + elif spec_signal == "has_speculative_draft_tokens": + attn_metadata.runtime_features.has_speculative_draft_tokens = True + + attn_metadata.request_ids = list(range(1, num_seqs + 1)) + attn_metadata.prompt_lens = prompt_lens + attn_metadata._seq_lens = torch.tensor(seq_lens_q, dtype=torch.int32) + attn_metadata._seq_lens_kv = torch.tensor(seq_lens_q, dtype=torch.int32) + attn_metadata.kv_cache_params = KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=num_cached_tokens_per_seq, + num_extra_kv_tokens=(7 if spec_signal == "num_extra_kv_tokens" else 0)) + + attn_metadata.prepare() + + if spec_signal is None: + # Non-speculative: capped at ceil(max kv len / tokens_per_block). + expected_max_blocks = -(-max(prompt_lens) // + mock_kv_cache_manager.tokens_per_block) + else: + expected_max_blocks = None + call_kwargs = mock_kv_cache_manager.copy_batch_block_offsets.call_args.kwargs + assert call_kwargs["max_blocks"] == expected_max_blocks + if mock_draft_manager is not None: + draft_kwargs = mock_draft_manager.copy_batch_block_offsets.call_args.kwargs + assert draft_kwargs["max_blocks"] is None + + @pytest.mark.parametrize( "use_cuda_graph,attn_backend,disable_overlap_scheduler,enable_block_reuse,use_one_model,enable_chunked_prefill,use_chain_drafter,multi_batch,attention_dp,use_hf_speculative_model", [