From 9a47c70e65498c3d5f8a41795f44f1291ad5bc10 Mon Sep 17 00:00:00 2001 From: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> Date: Wed, 14 Oct 2026 15:33:10 -0700 Subject: [PATCH 1/2] [TRTLLM-14620][feat] Add DSpark ragged graph and runtime integration Add exact-shape verification graph/runtime integration, isolate confidence graph policy, and preserve constructor-time ragged metadata opt-in. Signed-off-by: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> --- .../_torch/pyexecutor/cuda_graph_runner.py | 502 ++++++++++-- .../_torch/pyexecutor/engine/metadata.py | 11 + .../engine/runners/decoder/graphs.py | 68 +- .../engine/runners/decoder/inputs.py | 718 +++++++++++++++++- .../engine/runners/decoder/metadata.py | 138 ++++ .../engine/runners/decoder/runner.py | 75 +- .../engine/runners/decoder/warmup.py | 2 + .../_torch/pyexecutor/model_engine.py | 55 ++ tensorrt_llm/_torch/speculative/dspark.py | 13 +- .../_torch/speculative/dspark_select.py | 206 +++++ tensorrt_llm/_torch/speculative/interface.py | 46 ++ .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../test_phase_selective_forward.py | 2 + .../executor/test_moe_comm_timeout_warmup.py | 3 + .../executor/test_pytorch_model_engine.py | 11 + .../test_cuda_graph_dummy_identity.py | 94 +++ .../test_dspark_cuda_graph_runtime.py | 301 ++++++++ .../test_dspark_model_engine_runtime.py | 701 +++++++++++++++++ .../hw_agnostic/test_dspark_select.py | 71 ++ .../hw_agnostic/test_dspark_worker.py | 4 +- .../test_device_window_sampling_layout.py | 114 +++ 21 files changed, 3019 insertions(+), 117 deletions(-) create mode 100644 tensorrt_llm/_torch/speculative/dspark_select.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_cuda_graph_dummy_identity.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph_runtime.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_select.py create mode 100644 tests/unittest/_torch/speculative/test_device_window_sampling_layout.py diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 235dfdbb26cd..ba395157fd36 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -31,8 +31,8 @@ from ..speculative.utils import get_draft_kv_cache_manager from ..utils import make_weak_ref, piecewise_cuda_graph from .llm_request import LlmRequest, get_draft_token_length -from .resource_manager import (BaseResourceManager, ResourceManager, - ResourceManagerType) +from .resource_manager import (BaseResourceManager, NoFreeSlotsError, + ResourceManager, ResourceManagerType) from .sampler import SampleStateTensors from .sampler.sampler_common import SampleType from .scheduler import ScheduledRequests @@ -49,6 +49,32 @@ def _strict_buffer_check_enabled() -> bool: # A large prime number used for dummy request IDs to avoid collisions CUDA_GRAPH_DUMMY_REQUEST_ID = (1 << 64) - 1 + + +def cuda_graph_dummy_request_id(runtime_draft_len: int, *, variant: int, + max_draft_len: int) -> int: + """Return the disjoint padding-dummy ID for a draft/window variant.""" + runtime_draft_len = int(runtime_draft_len) + variant = int(variant) + max_draft_len = int(max_draft_len) + if (variant not in (0, 1) or max_draft_len < 0 or runtime_draft_len < 0 + or runtime_draft_len > max_draft_len): + raise ValueError("invalid CUDA padding dummy identity: " + f"draft_len={runtime_draft_len}, variant={variant}, " + f"max_draft_len={max_draft_len}") + return (CUDA_GRAPH_DUMMY_REQUEST_ID - variant * (max_draft_len + 1) - + runtime_draft_len) + + +def _padding_dummy_draft_token_capacity( + spec_config: Optional[DecodingBaseConfig], + allocated_draft_tokens: int) -> int: + """Keep DSpark's physical-K bound without restricting tree speculation.""" + if getattr(spec_config, "decoding_type", None) == "DSpark": + return int(spec_config.max_draft_len) + return int(allocated_draft_tokens) + + # Gen dummies get prompt_len = token_num - 1. Before capturing enc-dec decode # graphs, EncoderDecoderCudaGraphs.prepare_capture_batch temporarily runs each # dummy generation request as a one-token context chunk to write its cross-KV @@ -61,6 +87,45 @@ def get_mrope_dummy_seq_slot(max_num_tokens: int, pp_size: int) -> int: return max_num_tokens * pp_size +@dataclass(frozen=True) +class ADPShapeAgreement: + """One DSpark policy collective's queue and graph-shape agreement. + + The object is local process state, so ``batch_identity`` is intentionally + the Python identity on this rank. The collective-derived fields are equal + across ranks. Only a finalized exact-policy agreement may suppress the + later graph-shape and padding collectives. + """ + + iteration: int + batch_identity: int + local_batch_size: int + peer_batch_sizes: Tuple[int, ...] + all_can_graph: bool + widest_batch_size: int + graph_batch_size: int + draft_len: int + padding_ready: bool + ragged_bucket: Optional[int] = None + reuse_graph_shape: bool = False + + @property + def can_queue(self) -> bool: + return bool(self.peer_batch_sizes) and all(self.peer_batch_sizes) + + @property + def can_queue_this_rank(self) -> bool: + return self.local_batch_size > 0 + + def matches(self, batch: ScheduledRequests, iteration: int, *, + padded: bool) -> bool: + expected_size = (self.graph_batch_size + if padded else self.local_batch_size) + return (self.iteration == int(iteration) + and self.batch_identity == id(batch) + and batch.batch_size == expected_size) + + class KeyType(NamedTuple): batch_size: int draft_len: int @@ -78,6 +143,9 @@ class KeyType(NamedTuple): num_encoder_tokens: int = 0 peft_cache_data_type: Optional[torch.dtype] = None use_lora_graph: bool = False + # Total submitted verifier tokens for a ragged generation graph. ``None`` + # retains the existing uniform graph family. + ragged_verify_bucket: Optional[int] = None def _save_spec_decode_capture_state( @@ -451,7 +519,8 @@ def get_graph_key( context_query_len=context_query_len, num_encoder_tokens=num_encoder_tokens, peft_cache_data_type=peft_cache_data_type, - use_lora_graph=use_lora_graph) + use_lora_graph=use_lora_graph, + ragged_verify_bucket=self._ragged_verify_bucket(batch)) def set_capture_sample_type(self, sample_type: Optional[SampleType]) -> None: @@ -528,6 +597,32 @@ def _get_compatible_mixed_encoder_decoder_key(self, return min(compatible_keys, key=lambda captured_key: captured_key.num_encoder_tokens) + def _ragged_verify_bucket(self, batch: ScheduledRequests) -> Optional[int]: + """Optional token bucket supplied by a specialized graph runner.""" + return None + + def _agree_graph_shape( + self, batch: ScheduledRequests, can_run_cuda_graph: bool, + promoted_context_request_ids: frozenset[int] + ) -> Tuple[bool, Optional[SampleType]]: + """Agree on graph eligibility, rows and sampling tier across ADP ranks.""" + if not (self.enabled and self.config.enable_attention_dp + and self.config.mapping.tp_size > 1): + return True, None + local_sample_type = self._resolve_sample_type( + batch, promoted_context_request_ids) + graph_batch_info = self._gather_adp_graph_batch_info( + batch, can_run_cuda_graph, local_sample_type) + all_can_graph = all(rank_info[0] for rank_info in graph_batch_info) + all_batch_sizes_equal = all(rank_info[1] == graph_batch_info[0][1] + for rank_info in graph_batch_info) + if not all_can_graph or not all_batch_sizes_equal: + return False, None + forced_sample_type = (local_sample_type if all( + rank_info[2] == graph_batch_info[0][2] + for rank_info in graph_batch_info) else SampleType.FULL) + return True, forced_sample_type + @staticmethod def _get_mrope_position_delta(request: Any) -> Optional[Any]: mrope_position_delta = getattr(request, "py_mrope_position_delta", None) @@ -591,28 +686,10 @@ def maybe_get_cuda_graph( is_mixed_encoder_decoder = self._is_mixed_encoder_decoder_batch(batch) can_run_cuda_graph = self._can_run_cuda_graph_batch(batch) batch_size = batch.batch_size - # The sampling tier joins the graph key, so it has to agree across the - # attention-DP ranks or they would replay different graphs -- the same - # reason can_run_cuda_graph and batch_size are all-gathered here. Ranks - # see different batches (an idle one resolves FULL), so fall back to - # FULL for the whole group unless every rank picked the same tier. - forced_sample_type: Optional[SampleType] = None - if self.enabled and self.config.enable_attention_dp and self.config.mapping.tp_size > 1: - local_sample_type = self._resolve_sample_type( - batch, promoted_context_request_ids) - graph_batch_info = self._gather_adp_graph_batch_info( - batch, can_run_cuda_graph, local_sample_type) - all_can_run_cuda_graph = all(rank_info[0] - for rank_info in graph_batch_info) - all_batch_sizes_equal = all(rank_info[1] == graph_batch_info[0][1] - for rank_info in graph_batch_info) - - if not all_can_run_cuda_graph or not all_batch_sizes_equal: - return None, None, None - - forced_sample_type = (local_sample_type if all( - rank_info[2] == graph_batch_info[0][2] - for rank_info in graph_batch_info) else SampleType.FULL) + shapes_agree, forced_sample_type = self._agree_graph_shape( + batch, can_run_cuda_graph, promoted_context_request_ids) + if not shapes_agree: + return None, None, None if not self.enabled or not can_run_cuda_graph: return None, None, None @@ -1020,6 +1097,25 @@ def _check_strict_buffer_stability(self, stored_meta: Dict[str, self._assert_graph_tensor_ptrs_stable( stored_meta["spec_metadata"], stored_meta["spec_metadata_ptrs"]) + def will_pad_to(self, padded_bs: int, num_requests: int) -> bool: + """Whether a batch of ``num_requests`` can actually reach ``padded_bs``. + + Mirrors the SIZE guards in :meth:`_get_padded_batch` only; the + resource-dependent bails (dummy KV capacity, encoder-decoder state) + are deliberately not predicted. The ragged fit sizes its bucket grid + from the padded row count, so a wrong yes strands the fitted total in + no captured bucket. + """ + if padded_bs <= 0 or num_requests <= 0: + return False + if padded_bs not in self.supported_batch_sizes: + return False + if padded_bs == num_requests: + return True # already at a captured size, nothing to add + # _get_padded_batch bails when padding_size + batch_size exceeds the + # configured batch size; that sum is just padded_bs. + return padded_bs <= self.config.batch_size + def _get_padded_batch(self, batch: ScheduledRequests, resource_manager: ResourceManager, runtime_draft_len: int) -> int: @@ -1046,20 +1142,7 @@ def _get_padded_batch(self, batch: ScheduledRequests, """ can_run_cuda_graph = self._can_run_cuda_graph_batch(batch) batch_size = batch.batch_size - new_batch_size = batch_size - - if self.enabled and self.config.enable_attention_dp and self.config.mapping.tp_size > 1: - # Padding only agrees on eligibility and batch size; the tier slot - # is unused here. It is still sent so both per-iteration exchanges - # share one fixed-size payload, and a constant keeps this path from - # resolving a tier for a batch it is about to change. - graph_batch_info = self._gather_adp_graph_batch_info( - batch, can_run_cuda_graph, SampleType.FULL) - all_can_run_cuda_graph = all(rank_info[0] - for rank_info in graph_batch_info) - if all_can_run_cuda_graph: - new_batch_size = max(rank_info[1] - for rank_info in graph_batch_info) + new_batch_size = self._get_padding_batch_size(batch, can_run_cuda_graph) if (not self.enabled or not self.padding_enabled or not can_run_cuda_graph @@ -1072,6 +1155,7 @@ def _get_padded_batch(self, batch: ScheduledRequests, # otherwise this reduces to plain batch-size rounding. padded_batch_size = self._round_up_batch_size_with_draft_len( new_batch_size, runtime_draft_len) + self._validate_padding_batch_size(batch, padded_batch_size) if batch_size == padded_batch_size: return 0 @@ -1096,12 +1180,56 @@ def _get_padded_batch(self, batch: ScheduledRequests, key=f"cuda_graph_padding_dummy_fallback_{runtime_draft_len}") return 0 + return self._append_padding_requests(batch, resource_manager, + runtime_draft_len, + padding_dummy_request, + padding_size) + + def _get_padding_batch_size(self, batch: ScheduledRequests, + can_run_cuda_graph: bool) -> int: + new_batch_size = batch.batch_size + if (self.enabled and self.config.enable_attention_dp + and self.config.mapping.tp_size > 1): + graph_batch_info = self._gather_adp_graph_batch_info( + batch, can_run_cuda_graph, SampleType.FULL) + if all(rank_info[0] for rank_info in graph_batch_info): + new_batch_size = max(rank_info[1] + for rank_info in graph_batch_info) + return new_batch_size + + def _validate_padding_batch_size(self, batch: ScheduledRequests, + padded_batch_size: int) -> None: + """Validate a specialized runner's pre-agreed padding shape.""" + + def _append_padding_requests(self, batch: ScheduledRequests, + resource_manager: ResourceManager, + runtime_draft_len: int, + padding_dummy_request: LlmRequest, + padding_size: int) -> int: batch.generation_requests.extend([padding_dummy_request] * padding_size) return padding_size - def _get_or_create_padding_dummy( - self, resource_manager: ResourceManager, - runtime_draft_len: int) -> Optional[LlmRequest]: + def _padding_dummy_cache(self, variant: int) -> Dict[int, LlmRequest]: + if variant != 0: + raise ValueError( + f"unsupported CUDA padding dummy variant {variant}") + return self.padding_dummy_requests + + def _padding_dummy_caches(self) -> Tuple[Dict[int, LlmRequest], ...]: + return (self.padding_dummy_requests, ) + + def reset_padding_dummies(self) -> None: + """Drop cached request objects when their resource manager changes.""" + self.padding_dummy_requests = {} + + def _on_padding_dummy_release(self) -> None: + """Invalidate specialized padding agreements after resource release.""" + + def _get_or_create_padding_dummy(self, + resource_manager: ResourceManager, + runtime_draft_len: int, + *, + variant: int = 0) -> Optional[LlmRequest]: """Returns the padding dummy request for the given draft length, allocating it from the KV cache manager on first use. @@ -1109,8 +1237,9 @@ def _get_or_create_padding_dummy( (e.g. saturated cache); the batch cannot be padded until a retry succeeds. """ - if runtime_draft_len in self.padding_dummy_requests: - return self.padding_dummy_requests[runtime_draft_len] + cache = self._padding_dummy_cache(variant) + if runtime_draft_len in cache: + return cache[runtime_draft_len] kv_cache_manager = resource_manager.get_resource_manager( self.config.kv_cache_manager_key) @@ -1133,8 +1262,10 @@ def _get_or_create_padding_dummy( draft_kv_cache_manager = get_draft_kv_cache_manager( self.spec_config, resource_manager) - # Use unique dummy request ID per draft length - dummy_request_id = CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len + max_draft_len = _padding_dummy_draft_token_capacity( + self.spec_config, runtime_draft_token_buffer_width) + dummy_request_id = cuda_graph_dummy_request_id( + runtime_draft_len, variant=variant, max_draft_len=max_draft_len) dummy_request = kv_cache_manager.add_dummy_requests( [dummy_request_id], token_nums=[ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM] @@ -1151,17 +1282,43 @@ def _get_or_create_padding_dummy( return None dummy_request = dummy_request[0] dummy_request.is_cuda_graph_dummy = True - if self.is_encoder_decoder: - if not self._add_cross_dummy_request( - dummy_request, resource_manager, dummy_encoder_output_len, - draft_kv_cache_manager): - return None + registered_managers: List[BaseResourceManager] = [kv_cache_manager] + if (draft_kv_cache_manager is not None + and draft_kv_cache_manager is not kv_cache_manager): + registered_managers.append(draft_kv_cache_manager) - spec_res_mgr = resource_manager.get_resource_manager( - ResourceManagerType.SPEC_RESOURCE_MANAGER) - if spec_res_mgr: - spec_res_mgr.add_dummy_requests([dummy_request_id]) - self.padding_dummy_requests[runtime_draft_len] = dummy_request + def rollback() -> None: + for manager in reversed(registered_managers): + manager.free_resources(dummy_request) + + try: + if self.is_encoder_decoder: + cross_kv_cache_manager = resource_manager.get_resource_manager( + ResourceManagerType.CROSS_KV_CACHE_MANAGER) + if (cross_kv_cache_manager is None + or not self._add_cross_dummy_request( + dummy_request, resource_manager, + dummy_encoder_output_len, draft_kv_cache_manager)): + rollback() + return None + if not any(cross_kv_cache_manager is manager + for manager in registered_managers): + registered_managers.append(cross_kv_cache_manager) + + spec_res_mgr = resource_manager.get_resource_manager( + ResourceManagerType.SPEC_RESOURCE_MANAGER) + if spec_res_mgr: + spec_res_mgr.add_dummy_requests([dummy_request_id]) + if not any(spec_res_mgr is manager + for manager in registered_managers): + registered_managers.append(spec_res_mgr) + except NoFreeSlotsError: + rollback() + return None + except Exception: + rollback() + raise + cache[runtime_draft_len] = dummy_request return dummy_request def _padding_dummy_managers( @@ -1198,20 +1355,27 @@ def release_padding_dummy(self, resource_manager: ResourceManager, manager that allocated part of it, and drops it from the runner so a later padded step re-creates it. - One dummy request ID is spread across up to four managers -- the main - KV cache manager, the one-model draft KV cache manager, the + Each dummy request ID is spread across up to four managers -- + the main KV cache manager, the one-model draft KV cache manager, the speculative resource manager slot and the encoder-decoder cross-KV - cache manager. Releasing only the main one leaves the others holding - the ID, and re-creation reuses the same - ``CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len``. + cache manager. Releasing only the main one leaves the others holding + the ID; re-creation deterministically reuses the same variant ID. Returns True if a dummy was held for that draft length. """ - dummy_request = self.padding_dummy_requests.pop(runtime_draft_len, None) - if dummy_request is None: + dummy_requests = [ + cache.pop(runtime_draft_len, None) + for cache in self._padding_dummy_caches() + ] + dummy_requests = [ + request for request in dummy_requests if request is not None + ] + if not dummy_requests: return False - for manager in self._padding_dummy_managers(resource_manager): - manager.free_resources(dummy_request) + self._on_padding_dummy_release() + for dummy_request in dummy_requests: + for manager in self._padding_dummy_managers(resource_manager): + manager.free_resources(dummy_request) return True def _can_pad_any_batch(self, runtime_draft_len: int) -> bool: @@ -1270,6 +1434,12 @@ def preallocate_padding_dummies(self, logger.info( "Pre-allocated the CUDA graph padding dummy request " f"(draft_len={draft_len}) at warmup.") + self._preallocate_extra_padding_dummy(resource_manager, draft_len) + + def _preallocate_extra_padding_dummy(self, + resource_manager: ResourceManager, + draft_len: int) -> None: + """Reserve any additional padding variants required by a subclass.""" def _add_cross_dummy_request( self, dummy_request: LlmRequest, resource_manager: ResourceManager, @@ -1293,11 +1463,6 @@ def _add_cross_dummy_request( if cross_dummy_requests is not None: return True - kv_cache_manager = resource_manager.get_resource_manager( - self.config.kv_cache_manager_key) - kv_cache_manager.free_resources(dummy_request) - if draft_kv_cache_manager is not None: - draft_kv_cache_manager.free_resources(dummy_request) return False @staticmethod @@ -1368,7 +1533,7 @@ def clear(self, *, release_nccl_window_owners: bool = True): self.graph_outputs.clear() self.graph_metadata.clear() self._pre_replay_position_ids = None - self.padding_dummy_requests = {} + self.reset_padding_dummies() del self.memory_pool self.memory_pool = None if release_nccl_window_owners and graph_pool is not None: @@ -1376,6 +1541,201 @@ def clear(self, *, release_nccl_window_owners: bool = True): torch.cuda.empty_cache() +class DSparkCUDAGraphRunner(CUDAGraphRunner): + """Ragged verification graph policy for confidence-enabled DSpark.""" + + def __init__(self, config: CUDAGraphRunnerConfig) -> None: + super().__init__(config) + # Zero-real exact cells use low/high rows to represent V=G*q+r. + self.ragged_pad_verify_len = 0 + self.ragged_zero_real_high_rows = 0 + # The decoder publishes the fitted bucket before graph lookup. + self.agreed_ragged_bucket: Optional[int] = None + self.adp_shape_agreement: Optional[ADPShapeAgreement] = None + self.secondary_padding_dummy_requests: Dict[int, LlmRequest] = {} + + def _ragged_verify_bucket(self, batch: ScheduledRequests) -> Optional[int]: + """Return the agreed token bucket only when every row has a window.""" + # A window-less batch must not replay stale token-major row maps. + if any( + getattr(request, "py_verify_len", None) is None + for request in batch.generation_requests): + return None + # Use the globally fitted bucket, never a rank-local window sum. + return self.agreed_ragged_bucket + + @staticmethod + def _local_draft_len(batch: ScheduledRequests) -> int: + """Report malformed batches' maximum so the shape gate declines them.""" + return max((len(request.py_draft_tokens) + for request in batch.generation_requests), + default=0) + + def _agree_graph_shape( + self, batch: ScheduledRequests, can_run_cuda_graph: bool, + promoted_context_request_ids: frozenset[int] + ) -> Tuple[bool, Optional[SampleType]]: + agreement = self.adp_shape_agreement + # The policy agreement covers shape, not the in-graph sampling tier. + reuse_agreement = bool( + agreement is not None and agreement.reuse_graph_shape + and not self.enable_in_graph_sampling + and agreement.matches(batch, agreement.iteration, padded=True)) + if reuse_agreement: + if (not can_run_cuda_graph + or batch.batch_size != agreement.graph_batch_size or + self._ragged_verify_bucket(batch) != agreement.ragged_bucket + or self._local_draft_len(batch) != agreement.draft_len): + raise RuntimeError( + "DSpark cached ADP graph-shape agreement no longer " + "matches the local replay key") + adp_graph_agreement = (self.enabled and self.config.enable_attention_dp + and self.config.mapping.tp_size > 1) + local_sample_type = self._resolve_sample_type( + batch, + promoted_context_request_ids) if adp_graph_agreement else None + if not adp_graph_agreement or reuse_agreement: + return True, None + + # Replay freezes all_rank_num_tokens. Compare every rank-local shape + # component before launch; -1 denotes a non-applicable field. + local_bucket = self._ragged_verify_bucket( + batch) if can_run_cuda_graph else None + peer_shapes = self.config.dist.tp_allgather_int64([ + can_run_cuda_graph, + batch.batch_size, + -1 if local_bucket is None else int(local_bucket), + self._local_draft_len(batch) if can_run_cuda_graph else -1, + local_sample_type.value, + ]).tolist() + if (not all(peer[0] for peer in peer_shapes) + or not all(peer[1:4] == peer_shapes[0][1:4] + for peer in peer_shapes)): + return False, None + forced_sample_type = (local_sample_type if all( + peer[4] == peer_shapes[0][4] + for peer in peer_shapes) else SampleType.FULL) + return True, forced_sample_type + + def _get_num_tokens_for_key(self, key: KeyType) -> int: + if key.ragged_verify_bucket is not None: + if key.num_contexts: + raise ValueError( + "Ragged verifier graphs do not support mixed context batches." + ) + return int(key.ragged_verify_bucket) + return super()._get_num_tokens_for_key(key) + + def _reusable_padding_agreement( + self, batch: ScheduledRequests) -> Optional[ADPShapeAgreement]: + agreement = self.adp_shape_agreement + if (agreement is not None and agreement.reuse_graph_shape and + agreement.matches(batch, agreement.iteration, padded=False)): + return agreement + return None + + def _get_padding_batch_size(self, batch: ScheduledRequests, + can_run_cuda_graph: bool) -> int: + agreement = self._reusable_padding_agreement(batch) + if agreement is None: + return super()._get_padding_batch_size(batch, can_run_cuda_graph) + if not can_run_cuda_graph or not agreement.padding_ready: + raise RuntimeError( + "DSpark cached ADP padding agreement lost a local " + "eligibility invariant") + return agreement.widest_batch_size + + def _validate_padding_batch_size(self, batch: ScheduledRequests, + padded_batch_size: int) -> None: + agreement = self._reusable_padding_agreement(batch) + if (agreement is not None + and padded_batch_size != agreement.graph_batch_size): + raise RuntimeError( + "DSpark cached ADP padding agreement selected a different " + "graph batch size") + + def _append_padding_requests(self, batch: ScheduledRequests, + resource_manager: ResourceManager, + runtime_draft_len: int, + padding_dummy_request: LlmRequest, + padding_size: int) -> int: + # The ragged fit chooses windows before cached dummies are appended. + secondary_padding_dummy = None + high_padding_rows = 0 + if self.ragged_zero_real_high_rows: + if (batch.batch_size != 1 + or not batch.generation_requests[0].is_attention_dp_dummy + or self.ragged_zero_real_high_rows + > batch.batch_size + padding_size): + raise RuntimeError( + "DSpark zero-real heterogeneous padding lost its single " + "scheduled attention-DP dummy invariant") + high_padding_rows = self.ragged_zero_real_high_rows - 1 + if high_padding_rows > 0: + secondary_padding_dummy = self._get_or_create_padding_dummy( + resource_manager, runtime_draft_len, variant=1) + if secondary_padding_dummy is None: + logger.warning_once( + "Failed to allocate the secondary DSpark CUDA graph " + "padding dummy; falling back to native verification.", + key=("dspark_secondary_padding_dummy_fallback_" + f"{runtime_draft_len}")) + return 0 + ragged = any( + getattr(request, "py_verify_len", None) is not None + for request in batch.generation_requests) + padding_dummy_request.py_verify_len = (self.ragged_pad_verify_len + if ragged else None) + if secondary_padding_dummy is not None: + secondary_padding_dummy.py_verify_len = ( + self.ragged_pad_verify_len + 1 if ragged else None) + low_padding_rows = padding_size - high_padding_rows + if low_padding_rows < 0: + raise RuntimeError( + "DSpark zero-real high-row count exceeds CUDA padding rows") + if high_padding_rows: + batch.generation_requests.extend([secondary_padding_dummy] * + high_padding_rows) + batch.generation_requests.extend([padding_dummy_request] * + low_padding_rows) + return padding_size + + def _padding_dummy_cache(self, variant: int) -> Dict[int, LlmRequest]: + if variant == 1: + return self.secondary_padding_dummy_requests + return super()._padding_dummy_cache(variant) + + def _padding_dummy_caches(self) -> Tuple[Dict[int, LlmRequest], ...]: + return super()._padding_dummy_caches() + ( + self.secondary_padding_dummy_requests, ) + + def reset_padding_dummies(self) -> None: + super().reset_padding_dummies() + self.secondary_padding_dummy_requests = {} + self.ragged_zero_real_high_rows = 0 + self.adp_shape_agreement = None + + def _on_padding_dummy_release(self) -> None: + # Releasing a dummy invalidates the agreement that certified it. + self.adp_shape_agreement = None + + def _preallocate_extra_padding_dummy(self, + resource_manager: ResourceManager, + draft_len: int) -> None: + if not getattr(self.spec_config, "confidence_sps_table_path", None): + return + if self._get_or_create_padding_dummy(resource_manager, + draft_len, + variant=1) is None: + logger.warning("Could not pre-allocate the secondary DSpark CUDA " + "graph padding dummy " + f"(draft_len={draft_len}); non-divisible zero-real " + "exact cells will remain native.") + else: + logger.info("Pre-allocated the secondary DSpark CUDA graph " + f"padding dummy (draft_len={draft_len}).") + + EncoderKeyType: TypeAlias = Tuple[int, int, int] _ENCODER_SOURCE_SEQ_LENS = "_encoder_source_seq_lens" _ENCODER_SOURCE_TO_SLOT = "_encoder_source_to_slot" diff --git a/tensorrt_llm/_torch/pyexecutor/engine/metadata.py b/tensorrt_llm/_torch/pyexecutor/engine/metadata.py index 8b130ab9877c..e08b1e1dccc2 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/metadata.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/metadata.py @@ -59,6 +59,7 @@ def build_attention_metadata( draft_kv_cache_manager: KVCacheManager | KVCacheManagerV2 | None = None, enable_context_mla_with_cached_kv: bool | None = None, max_num_heads_per_kv: int | None = None, + enable_ragged_verification: bool = False, ) -> AttentionMetadata: """Construct attention metadata and resolve model-derived inputs.""" pretrained_config = model_config.pretrained_config @@ -75,6 +76,15 @@ def build_attention_metadata( if sparse_attention_config is not None else None ) + metadata_kwargs = {} + if enable_ragged_verification: + from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import ( + DSAtrtllmAttentionMetadata, + ) + + if not issubclass(attention_backend.Metadata, DSAtrtllmAttentionMetadata): + raise ValueError("Ragged verification requires the DSA target metadata backend") + metadata_kwargs["enable_ragged_verification"] = True return attention_backend.Metadata( max_num_requests=max_batch_size, max_num_tokens=max_num_tokens, @@ -88,4 +98,5 @@ def build_attention_metadata( cache_indirection=cache_indirection, max_num_heads_per_kv=max_num_heads_per_kv, sparse_metadata_params=sparse_metadata_params, + **metadata_kwargs, ) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py index 2f3b6482b5f0..5fe5c5d70e54 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py @@ -58,6 +58,7 @@ from ..common import get_top_level_model, prepare_multimodal_indices from .config import DecoderRunnerConfig from .dummy_requests import DummyRequestFactory +from .inputs import DecoderInputPreparer from .metadata import DecoderMetadata from .warmup import WarmupAgreement @@ -205,8 +206,10 @@ def __init__( torch_compile_bypass: Callable[[], contextlib.AbstractContextManager[None]], context_promotion_supported: bool, warmup_inputs: torch.Tensor | None = None, + inputs: DecoderInputPreparer | None = None, ) -> None: self.cuda_graph_runner = cuda_graph_runner + self._inputs = inputs self.breakable_cuda_graph_runner = breakable_cuda_graph_runner self._config = config self.mapping = mapping @@ -329,7 +332,7 @@ def captures_model_body(self, key: KeyType | None) -> bool: def reset_padding_dummies(self) -> None: """Drop padding dummies that belong to an earlier KV cache manager.""" - self.cuda_graph_runner.padding_dummy_requests = {} + self.cuda_graph_runner.reset_padding_dummies() def preallocate_padding_dummies(self, resource_manager: ResourceManager) -> None: """Allocate the padding dummies while the KV cache is still empty.""" @@ -602,7 +605,9 @@ def capture_model(model_inputs: dict[str, Any]) -> Any: restore_attn_metadata_after_draft_replay(attn_metadata, saved_draft) return outputs - def _get_graphs_to_capture(self, cuda_graph_batch_sizes: list[int]) -> list[tuple[int, int]]: + def _get_graphs_to_capture( + self, cuda_graph_batch_sizes: list[int] + ) -> list[tuple[int, int] | tuple[int, int, int]]: """Determine which (batch_size, draft_len) graphs to capture. Returns: @@ -624,6 +629,33 @@ def _get_graphs_to_capture(self, cuda_graph_batch_sizes: list[int]) -> list[tupl ) return graphs + if self._inputs is not None and self._inputs._dspark_confidence_enabled: + max_draft_len = int(self._config.spec_config.max_draft_len) + native_graphs = [(bs, max_draft_len) for bs in cuda_graph_batch_sizes] + if not self._inputs._dspark_trims_submitted_tokens: + return native_graphs + from tensorrt_llm._torch.speculative.dspark_planner import ExactSpsCostTable + + exact_table = self._inputs._dspark_sps_cost_table + if not isinstance(exact_table, ExactSpsCostTable): + raise RuntimeError( + "DSpark confidence graph capture requires an authenticated exact SPS table" + ) + ragged_graphs = [ + (bs, max_draft_len, verifier_budget) + for bs in cuda_graph_batch_sizes + for verifier_budget in exact_table.production_candidate_budgets(bs) + ] + # Full K keeps the ordinary native graph. Compact layouts have + # distinct keys; they cannot alias its uniform attention layout. + graphs = native_graphs + ragged_graphs + logger.info( + f"DSpark ragged verification: capturing {len(graphs)} graphs " + f"({len(native_graphs)} native static-K{max_draft_len} graphs + " + f"{len(ragged_graphs)} exact measured (G,V) cells)." + ) + return graphs + # Case 2: Static draft length # Match the runtime_draft_len semantics enforced in resolve_runtime_draft_len: # logical K for linear-tree modes, total tree tokens for tree decoding. @@ -758,7 +790,9 @@ def _run_capture_pass( pinned_tier = sample_type or SampleType.FULL self.cuda_graph_runner.set_capture_sample_type(pinned_tier) try: - for bs, draft_len in graphs_to_capture: + for entry in graphs_to_capture: + bs, draft_len = entry[:2] + verify_bucket = entry[2] if len(entry) == 3 else None if bs > self._config.max_batch_size: continue @@ -773,6 +807,8 @@ def _run_capture_pass( with self._dummy_requests.release_batch_context( warmup_request, resource_manager ) as batch: + if batch is not None and verify_bucket is not None: + self._set_warmup_ragged_windows(batch, verify_bucket, draft_len) if batch is None: # No KV cache space for this batch size. During KV # cache estimation this makes the profiling peak @@ -814,6 +850,8 @@ def _run_capture_pass( finally: self._force_lora_graph_for_capture = None self.cuda_graph_runner.set_capture_sample_type(None) + if self._inputs is not None and self._inputs._dspark_confidence_enabled: + self.cuda_graph_runner.agreed_ragged_bucket = None if self._lora.cuda_graph_manager is None: lora_graph_cases = [False] @@ -1009,3 +1047,27 @@ def release(self, *, release_nccl_window_owners: bool = True) -> None: self.breakable_cuda_graph_runner.clear( release_nccl_window_owners=release_nccl_window_owners ) + + def _set_warmup_ragged_windows(self, batch, verify_bucket: int, draft_len: int) -> None: + """Give a capture batch per-request windows summing to ``verify_bucket``. + + The warmup batch must hit the bucket total exactly, or the graph is + keyed under a total no runtime batch produces and every ragged step + silently falls back to eager. Any split with the right sum captures + the right shape (the raggedness lives in the contents of qo_indptr). + """ + requests = batch.generation_requests + if not requests: + return + from tensorrt_llm._torch.speculative.ragged_helpers import RaggedVerifyLayout + + lens = torch.ones(len(requests), dtype=torch.int32) + layout = RaggedVerifyLayout.from_verify_lens( + lens, graph_num_tokens=int(verify_bucket), total_verify_tokens=len(requests) + ) + filled = layout.fill_bucket(max_verify_len=1 + int(draft_len)) + for request, tokens in zip(requests, filled.verify_lens.tolist()): + request.py_verify_len = int(tokens) - 1 + # Publish to the runner too: the graph key carries this value, and the + # runtime fit never runs during capture. + self.cuda_graph_runner.agreed_ragged_bucket = int(verify_bucket) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py index 578ab3e68667..292344be85a6 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py @@ -21,7 +21,11 @@ from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import get_mrope_dummy_seq_slot from tensorrt_llm._torch.pyexecutor.guided_decoder import CapturableGuidedDecoder from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, get_draft_token_length +from tensorrt_llm._torch.pyexecutor.llm_request import ( + LlmRequest, + get_draft_token_length, + get_request_tokens_per_gen_step, +) from tensorrt_llm._torch.pyexecutor.resource_manager import ( KVCacheManager, ResourceManager, @@ -33,6 +37,7 @@ from tensorrt_llm._torch.speculative import SpecMetadata, get_num_extra_kv_tokens from tensorrt_llm._torch.speculative.dspark import DSparkSpecMetadata from tensorrt_llm._torch.speculative.interface import INVALID_PROMPT_LOOKAHEAD_TOKEN +from tensorrt_llm._torch.speculative.ragged_helpers import ragged_gather_index_lists from tensorrt_llm._torch.speculative.spec_sampler_base import SampleStateTensorsSpec from tensorrt_llm._torch.speculative.utils import get_static_draft_len from tensorrt_llm._torch.utils import helix_local_len, set_per_request_prefill_cuda_graph_flag @@ -135,6 +140,8 @@ class DecoderInputBuffers(InputBuffers): """ previous_batch_indices_cuda: torch.Tensor + ragged_verify_lens_cuda: torch.Tensor + ragged_qo_indptr_cuda: torch.Tensor # Host counter of the steady-state generation prepare cache. steady_gen_positions_pinned: torch.Tensor previous_pos_indices_cuda: torch.Tensor | None = None @@ -165,6 +172,12 @@ def for_config( previous_batch_indices_cuda=torch.empty( (config.max_num_tokens,), dtype=torch.int, device="cuda" ), + ragged_verify_lens_cuda=torch.empty( + (config.max_batch_size + 1,), dtype=torch.int, device="cuda" + ), + ragged_qo_indptr_cuda=torch.empty( + (config.max_batch_size + 2,), dtype=torch.int, device="cuda" + ), steady_gen_positions_pinned=torch.empty( (config.max_num_tokens,), dtype=torch.int, pin_memory=prefer_pinned() ), @@ -295,6 +308,9 @@ class HostInputs: num_tokens: int = 0 total_num_tokens: int = 0 previous_batch_len: int = 0 + runtime_draft_len: int = 0 + previous_batch_token_start: int | None = None + previous_batch_tokens_per_request: list[int] | None = None host_position_ids: torch.Tensor | None = None num_cached_tokens_snapshot: list[int] | None = None remote_tail_modes: set[Any] = field(default_factory=set) @@ -325,6 +341,8 @@ def __init__( mm_item_scheduler: MultimodalItemScheduler | None = None, ) -> None: self._buffers = buffers + self._pinned_host_cache = {} + self._pinned_host_active = {} self._mm_item_scheduler = mm_item_scheduler self._config = config self.mapping = mapping @@ -363,6 +381,38 @@ def __init__( os.getenv("TLLM_LOG_CACHED_KV_TOKENS_PER_REQ", "0") == "1" ) + spec_config = self._config.spec_config + self._dspark_confidence_enabled = bool( + spec_config is not None and getattr(spec_config, "enable_confidence_scheduling", False) + ) + self._dspark_trims_submitted_tokens = self._dspark_confidence_enabled + self._dspark_device_windows = bool( + self._dspark_trims_submitted_tokens + and getattr(spec_config, "enable_fused_confidence_scheduler", False) + ) + self._dspark_sps_cost_table = None + self._dspark_exact_candidate_cells = () + self._dspark_exact_identity_words = (0,) * 8 + table_path = getattr(spec_config, "confidence_sps_table_path", None) + if self._dspark_confidence_enabled and table_path: + from tensorrt_llm._torch.speculative.dspark_planner import load_runtime_sps_cost_table + + self._dspark_sps_cost_table, _ = load_runtime_sps_cost_table( + table_path, + graph_batch_sizes=self._config.cuda_graph_batch_sizes, + max_draft_len=int(spec_config.max_draft_len), + live_engine_fingerprint_path=getattr( + spec_config, "confidence_sps_live_fingerprint_path", None + ), + ) + self._dspark_exact_candidate_cells = self._dspark_sps_cost_table.candidate_cells() + self._dspark_exact_identity_words = ( + self._dspark_sps_cost_table.collective_identity_words + ) + self.previous_request_ids = [] + self.previous_verify_lens = [] if self._dspark_trims_submitted_tokens else None + self.has_previous_device_draft = False + def _new_extra_inputs(self) -> ExtraInputsCollector: """Return the collector of this pass's model inputs beyond the decoder's own.""" return _NO_EXTRA_INPUTS @@ -851,6 +901,7 @@ def prepare_inputs( if enable_spec_decode: self._gather_ids_buffer.update(torch.tensor(host.lists.gather_ids, dtype=torch.int)) + host.runtime_draft_len = runtime_draft_len self._prepare_attn_metadata( host, scheduled_requests, @@ -896,6 +947,13 @@ def prepare_inputs( ) self._record_iter_states(host, scheduled_requests) if not is_dummy: + self.previous_request_ids = host.lists.all_gen_request_ids + if self._dspark_trims_submitted_tokens: + self.previous_verify_lens = [ + getattr(request, "py_verify_len", None) + for request in scheduled_requests.generation_requests + ] + self.has_previous_device_draft = next_draft_tokens_device is not None self._record_steady_gen_cache( host, kv_cache_manager, attn_metadata, spec_metadata, new_tokens_device, lora_params ) @@ -1312,7 +1370,9 @@ def _pack_extend_requests( mrope_dummy_seq_slot = host.mrope_dummy_seq_slot has_gen_mrope_delta = host.has_gen_mrope_delta - runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 + previous_batch_token_start = None + previous_batch_tokens_per_request = [] if self._dspark_trims_submitted_tokens else None + extend_tokens_emitted = 0 # Under extend_ctx the attention kernel runs extend requests as chunked # context rows: their rotary comes from position_ids alone, and they take # no entry in the generation-indexed MRoPE delta read slots. @@ -1339,6 +1399,11 @@ def _pack_extend_requests( or request.is_dummy or request.py_batch_idx is None ): + num_draft_tokens = get_draft_token_length(request) + if self._dspark_trims_submitted_tokens: + num_draft_tokens = ( + get_request_tokens_per_gen_step(request, 1 + num_draft_tokens) - 1 + ) # get token ids, including input token ids and draft token ids. For these dummy requests, # no need to copy the token ids. if not (request.is_attention_dp_dummy or request.is_cuda_graph_dummy): @@ -1346,15 +1411,22 @@ def _pack_extend_requests( input_ids.append(request.get_tokens(0)[request.context_current_position]) else: input_ids.append(request.get_last_tokens(0)) - input_ids.extend(request.py_draft_tokens) - draft_tokens.extend(request.py_draft_tokens) + submitted_draft_tokens = ( + request.py_draft_tokens[:num_draft_tokens] + if self._dspark_trims_submitted_tokens + else request.py_draft_tokens + ) + input_ids.extend(submitted_draft_tokens) + draft_tokens.extend(submitted_draft_tokens) # get other ids and lengths - num_draft_tokens = get_draft_token_length(request) past_seen_token_num = ( request.context_current_position if is_promoted_context else request.max_beam_num_tokens - 1 ) + if self._dspark_trims_submitted_tokens: + request.cached_tokens = past_seen_token_num + extend_tokens_emitted += 1 + num_draft_tokens draft_lens.append(num_draft_tokens) if ( enable_spec_decode @@ -1406,29 +1478,35 @@ def _pack_extend_requests( previous_batch_idx = request.py_batch_idx request.py_batch_idx = request.py_seq_slot - sequence_lengths.append(runtime_tokens_per_gen_step) + request_tokens_per_step = runtime_tokens_per_gen_step + if self._dspark_trims_submitted_tokens: + request_tokens_per_step = get_request_tokens_per_gen_step( + request, runtime_tokens_per_gen_step + ) + if previous_batch_token_start is None: + previous_batch_token_start = extend_tokens_emitted + if previous_batch_tokens_per_request is not None: + previous_batch_tokens_per_request.append(request_tokens_per_step) + extend_tokens_emitted += request_tokens_per_step + sequence_lengths.append(request_tokens_per_step) past_seen_token_num = request.max_beam_num_tokens - 1 - draft_lens.append(runtime_draft_token_buffer_width) + draft_lens.append(request_tokens_per_step - 1) gather_ids.extend( - list(range(len(position_ids), len(position_ids) + runtime_tokens_per_gen_step)) + list(range(len(position_ids), len(position_ids) + request_tokens_per_step)) ) position_ids.extend( - list( - range( - past_seen_token_num, past_seen_token_num + runtime_tokens_per_gen_step - ) - ) + list(range(past_seen_token_num, past_seen_token_num + request_tokens_per_step)) ) # previous tensor previous_batch_indices.append(previous_batch_idx) - previous_pos_indices.extend([previous_batch_idx] * runtime_tokens_per_gen_step) + previous_pos_indices.extend([previous_batch_idx] * request_tokens_per_step) num_cached_tokens_per_seq.append( - past_seen_token_num - + runtime_tokens_per_gen_step - - request.py_num_compressed_tokens + past_seen_token_num + request_tokens_per_step - request.py_num_compressed_tokens ) + if self._dspark_trims_submitted_tokens: + request.cached_tokens = past_seen_token_num + request_tokens_per_step if _has_cp_helix: # In-flight predecessor: mirror the non-helix convention # above -- positions are packed from the stale base (the @@ -1437,7 +1515,7 @@ def _pack_extend_requests( # in recompute_helix_spec_buffers overrides them). The # base is reconstructed GLOBALLY (see the no-previous # branch: the token list is rank-local under helix). - group = runtime_tokens_per_gen_step + group = request_tokens_per_step base = _helix_pack_extend(request, group) local_full = _helix_local_len_host(base + group) helix_owned_new_tokens.append(0) @@ -1447,7 +1525,7 @@ def _pack_extend_requests( and spec_config.spec_dec_mode.extend_ctx(self._config.attention_backend) and spec_config.is_linear_tree ): - prompt_lengths.append(runtime_tokens_per_gen_step) + prompt_lengths.append(request_tokens_per_step) else: prompt_lengths.append(request.py_prompt_len) @@ -1505,6 +1583,8 @@ def _pack_extend_requests( extra_inputs.add_generation_request(request) host.has_gen_mrope_delta = has_gen_mrope_delta + host.previous_batch_token_start = previous_batch_token_start + host.previous_batch_tokens_per_request = previous_batch_tokens_per_request def _pack_generation_requests( self, @@ -1774,6 +1854,9 @@ def _copy_token_ids_to_device( extend_dummy_requests = host.lists.extend_dummy_requests previous_batch_len = len(previous_batch_indices) + self._dspark_prev_covers_batch = previous_batch_len > 0 and previous_batch_len == len( + extend_requests + ) - len(extend_dummy_requests) def previous_seq_slots_device() -> torch.Tensor: previous_slots = self._buffers.previous_batch_indices_cuda[:previous_batch_len] @@ -1816,7 +1899,8 @@ def previous_seq_slots_device() -> torch.Tensor: # below overwrite all rows _preprocess_inputs reads, so the buffers # need no zeroing first. if not ( - attn_metadata.is_cuda_graph + not self._dspark_trims_submitted_tokens + and attn_metadata.is_cuda_graph and previous_batch_len == num_extend_reqeust_wo_dummy and not extend_dummy_requests and not host.lists.generation_requests @@ -1827,7 +1911,11 @@ def previous_seq_slots_device() -> torch.Tensor: if previous_batch_len > 0: previous_slots = previous_seq_slots_device() - previous_batch_tokens = previous_batch_len * runtime_tokens_per_gen_step + is_ragged_gen = host.previous_batch_tokens_per_request is not None and any( + tokens != runtime_tokens_per_gen_step + for tokens in host.previous_batch_tokens_per_request + ) + previous_batch_tokens = len(previous_pos_indices) if previous_pos_indices != self._staged_previous_pos_indices: previous_pos_indices_host = torch.tensor( previous_pos_indices, dtype=torch.int, pin_memory=prefer_pinned() @@ -1857,7 +1945,12 @@ def previous_seq_slots_device() -> torch.Tensor: # For 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp. # Zeroed above. previous_begin = num_extend_reqeust_wo_dummy - previous_batch_len - if step_gather is None or not step_gather.gather( + pos_offsets_start = ( + host.previous_batch_token_start + if host.previous_batch_token_start is not None + else previous_begin * runtime_tokens_per_gen_step + ) + if is_ragged_gen or step_gather is None or not step_gather.gather( new_tokens_device, next_draft_tokens_device, new_tokens_lens_device, @@ -1871,23 +1964,44 @@ def previous_seq_slots_device() -> torch.Tensor: self._buffers.draft_tokens_cuda, num_draft_tokens, self._buffers.previous_pos_id_offsets_cuda, - previous_begin * runtime_tokens_per_gen_step, + pos_offsets_start, self._buffers.previous_kv_lens_offsets_cuda, previous_begin, ): # previous input ids - new_tokens = new_tokens_device.transpose(0, 1)[ - previous_slots, :runtime_tokens_per_gen_step - ].flatten() + if is_ragged_gen: + token_rows, token_cols = self._ragged_gather_indices( + previous_batch_indices, host.previous_batch_tokens_per_request + ) + new_tokens = new_tokens_device.transpose(0, 1)[ + token_rows, token_cols + ].flatten() + else: + new_tokens = new_tokens_device.transpose(0, 1)[ + previous_slots, :runtime_tokens_per_gen_step + ].flatten() self._buffers.input_ids_cuda[ num_tokens : num_tokens + previous_batch_tokens ].copy_(new_tokens, non_blocking=True) # previous draft tokens - previous_batch_draft_tokens = ( - previous_batch_len * runtime_draft_token_buffer_width - ) - if runtime_draft_token_buffer_width > 0: + if is_ragged_gen: + previous_batch_draft_tokens = previous_batch_tokens - previous_batch_len + if previous_batch_draft_tokens > 0: + draft_rows, draft_cols = self._ragged_gather_indices( + previous_batch_indices, + [tokens - 1 for tokens in host.previous_batch_tokens_per_request], + ) + self._buffers.draft_tokens_cuda[ + num_draft_tokens : num_draft_tokens + previous_batch_draft_tokens + ].copy_( + next_draft_tokens_device[draft_rows, draft_cols].flatten(), + non_blocking=True, + ) + elif runtime_draft_token_buffer_width > 0: + previous_batch_draft_tokens = ( + previous_batch_len * runtime_draft_token_buffer_width + ) self._buffers.draft_tokens_cuda[ num_draft_tokens : num_draft_tokens + previous_batch_draft_tokens ].copy_( @@ -1897,10 +2011,20 @@ def previous_seq_slots_device() -> torch.Tensor: non_blocking=True, ) # prepare data for the preprocess inputs - kv_len_offsets_device = new_tokens_lens_device - runtime_tokens_per_gen_step + if is_ragged_gen: + tokens_per_previous_request = torch.tensor( + host.previous_batch_tokens_per_request, + dtype=torch.long, + pin_memory=prefer_pinned(), + ).to(new_tokens_lens_device.device, non_blocking=True) + previous_kv_len_offsets = ( + new_tokens_lens_device[previous_slots] - tokens_per_previous_request + ).clamp_(min=-(runtime_draft_len + 1), max=runtime_draft_len + 1) + else: + kv_len_offsets_device = new_tokens_lens_device - runtime_tokens_per_gen_step + previous_kv_len_offsets = kv_len_offsets_device[previous_slots] self._buffers.previous_pos_id_offsets_cuda[ - previous_begin * runtime_tokens_per_gen_step : num_extend_reqeust_wo_dummy - * runtime_tokens_per_gen_step + pos_offsets_start : pos_offsets_start + previous_batch_tokens ].copy_( new_tokens_lens_device[ self._buffers.previous_pos_indices_cuda[0:previous_batch_tokens] @@ -1909,7 +2033,7 @@ def previous_seq_slots_device() -> torch.Tensor: ) self._buffers.previous_kv_lens_offsets_cuda[ previous_begin:num_extend_reqeust_wo_dummy - ].copy_(kv_len_offsets_device[previous_slots], non_blocking=True) + ].copy_(previous_kv_len_offsets, non_blocking=True) elif new_tokens_device is not None: seq_slots_device = previous_seq_slots_device() @@ -2025,7 +2149,21 @@ def _prepare_attn_metadata( ), ) - if not attn_metadata.is_cuda_graph: + if self._dspark_trims_submitted_tokens: + self._publish_gen_token_layout( + attn_metadata, + scheduled_requests.generation_requests, + runtime_draft_len=host.runtime_draft_len, + ) + refresh_seq_lens = not attn_metadata.is_cuda_graph + if ( + not refresh_seq_lens + and getattr(attn_metadata, "is_ragged_verify", False) + and attn_metadata.seq_lens is not None + and len(host.lists.sequence_lengths) == attn_metadata.seq_lens.shape[0] + ): + refresh_seq_lens = True + if refresh_seq_lens: # Assumes seq lens do not change between CUDA graph invocations. This applies # to draft sequences too. This means that all draft sequences must be padded. attn_metadata.seq_lens = torch.tensor( @@ -2232,6 +2370,10 @@ def _build_inputs( spec_metadata.host_position_ids = host.host_position_ids if host.context_prompt_lookahead is not None: spec_metadata.populate_context_prompt_lookahead(host.context_prompt_lookahead) + if self._dspark_trims_submitted_tokens: + self._attach_ragged_verify_layout( + spec_metadata, attn_metadata, scheduled_requests.generation_requests + ) # No-op for non 1-model spec_metadata.populate_sampling_params_for_one_model(scheduled_requests.all_requests()) spec_metadata.prepare() @@ -2515,3 +2657,509 @@ def postprocess_inputs( inputs["attn_metadata"].helix_position_offsets[:previous_batch_tokens] -= ( self._buffers.previous_pos_id_offsets_cuda[:previous_batch_tokens] ) + + def ragged_verify_token_buckets(self, padded_bs: int) -> list[int]: + """Measured compact ``V`` cells captured for one exact ``G``.""" + if self._config.spec_config is None: + return [] + from tensorrt_llm._torch.speculative.dspark_planner import ExactSpsCostTable + + exact_table = getattr(self, "_dspark_sps_cost_table", None) + if not isinstance(exact_table, ExactSpsCostTable): + raise RuntimeError( + "DSpark ragged verification requires an authenticated exact SPS table" + ) + return list(exact_table.production_candidate_budgets(int(padded_bs))) + + def fit_ragged_verify_lens( + self, + generation_requests, + verify_lens: list[int], + exact_shape: tuple[int, int, int], + peer_stats: list[list[int]] | None = None, + exact_zero_real: bool = False, + ) -> int | None: + """Validate and publish a globally selected exact ``(G,V)`` shape. + + ``peer_stats`` is every attention-DP rank's ``[num_requests, + total_tokens, all_can_graph]``, including this rank's: the bucket must + be sized from the cross-rank maximum or ranks pick different shapes, + and ``all_can_graph`` is the group's answer to whether that maximum + will be taken at all. + + ``exact_shape`` is the agreed ``(G,V,pad_tokens)`` from the exact SPS + selector. Changing V here would execute a different measured cell from + the one the policy priced, so this method only validates the split. + + ``exact_zero_real`` means this rank contributes no logical requests + but carries the one scheduler-owned attention-DP dummy. The dummy is + mapped onto the high row when ``V % G != 0``; CUDA padding supplies + the remaining low/high rows without publishing output or KV ownership + for a real request. + """ + runner = self.cuda_graph_runner + # Cleared on entry so every early return leaves no stale value: a + # bucket carried over from a previous step would key into the wrong + # graph. The padded row count follows the same rule. + runner.agreed_ragged_bucket = None + runner.ragged_zero_real_high_rows = 0 + self._dspark_last_padded_bs = None + if not runner.enabled or not verify_lens: + return None + if len(verify_lens) != len(generation_requests): + # A partially windowed batch would go ragged in the token layout + # while the spec metadata stayed uniform; refuse instead. + logger.debug( + f"DSpark ragged: got {len(verify_lens)} verify lengths for " + f"{len(generation_requests)} generation requests; falling back " + f"to uniform scheduling" + ) + return None + max_verify_len = 1 + int(self._config.spec_config.max_draft_len) + token_lens = [1 + int(v) for v in verify_lens] + + from tensorrt_llm._torch.speculative.ragged_helpers import resolve_ragged_pad_split + + # Resolve rows first, then tokens: the bucket grid depends on the + # widest rank's row count. `all_can_graph` (third peer-stat element; + # absent means single-rank) says whether _get_padded_batch will take + # the cross-rank row maximum at all -- only when EVERY rank can graph. + all_can_graph = bool(peer_stats[0][2]) if peer_stats and len(peer_stats[0]) > 2 else True + if not all_can_graph: + # No rank will replay a graph this step, and _get_padded_batch + # pads differently than this fit assumes on such steps, so the + # budget would go to rows that never appear. Decline (conservative: + # an eager ragged step would still save compute). + logger.debug( + "DSpark ragged: some rank cannot run a CUDA graph this step, " + "so no captured bucket applies; falling back to uniform" + ) + return None + widest_rows = ( + max([int(s[0]) for s in peer_stats], default=len(token_lens)) + if peer_stats + else len(token_lens) + ) + padded_bs = int(exact_shape[0]) + expected_padded_bs = runner._round_up_batch_size(widest_rows) + if padded_bs != expected_padded_bs: + logger.warning( + "DSpark exact ragged shape disagrees with graph ladder: " + f"selected G={padded_bs}, expected G={expected_padded_bs}" + ) + return None + if padded_bs == 0: + return None + # The fit assumes the batch will actually be padded to `padded_bs`; + # padding may decline, and then the fitted total lands in no captured + # bucket. Decline to go ragged rather than fit against rows that will + # not exist (`will_pad_to` mirrors only the SIZE guards; the group's + # graph answer is `all_can_graph` above). + if not runner.will_pad_to(padded_bs, len(token_lens)): + logger.debug( + f"DSpark ragged: padding to {padded_bs} rows is not available " + f"for {len(token_lens)} requests, so the bucket grid derived " + f"from it would not be realised; falling back to uniform" + ) + return None + buckets = self.ragged_verify_token_buckets(padded_bs) + if not buckets: + return None + bucket = int(exact_shape[1]) + if bucket not in buckets: + logger.warning(f"DSpark exact V={bucket} was not captured for G={padded_bs}") + return None + + full_bucket = int(padded_bs) * int(max_verify_len) + if bucket == full_bucket: + # A full verifier budget has no ragged work reduction. Preserve + # the native static-K path and its ordinary multi-G tail handling. + return None + + if exact_zero_real: + if ( + len(generation_requests) != 1 + or not generation_requests[0].is_attention_dp_dummy + or len(verify_lens) != 1 + ): + raise RuntimeError( + "DSpark zero-real exact fit lost its single scheduled " + "attention-DP dummy invariant" + ) + quotient, remainder = divmod(int(bucket), int(padded_bs)) + if quotient < 1 or quotient + int(remainder > 0) > max_verify_len: + raise RuntimeError( + "DSpark zero-real exact bucket has no bounded low/high " + f"row split: G={padded_bs}, V={bucket}, " + f"max_tokens={max_verify_len}" + ) + expected_scheduled_window = quotient - 1 + int(remainder > 0) + if int(verify_lens[0]) != expected_scheduled_window: + raise RuntimeError( + "DSpark zero-real scheduled dummy window differs from " + f"the exact quotient/remainder split: got={verify_lens[0]}, " + f"expected={expected_scheduled_window}" + ) + if remainder > 1: + draft_len = int(self._config.spec_config.max_draft_len) + if draft_len not in runner.secondary_padding_dummy_requests: + raise RuntimeError( + "DSpark secondary padding dummy disappeared after " + "the all-rank exact-cell agreement" + ) + self._dspark_last_padded_bs = int(padded_bs) + self._dspark_last_num_real = 0 + runner.agreed_ragged_bucket = int(bucket) + runner.ragged_pad_verify_len = int(quotient) - 1 + runner.ragged_zero_real_high_rows = int(remainder) + generation_requests[0].py_verify_len = expected_scheduled_window + return int(bucket) + + # Pad rows are appended later as a *single shared dummy object*, so + # they all carry the same window and their contribution must be + # decided here, not left to the fill. They take the minimum (one + # token) so real requests get the slack, growing only when the real + # rows cannot absorb the bucket. + n_real = len(token_lens) + n_pad = padded_bs - n_real + floor_tokens = sum(token_lens) + fixed_pad_len = int(exact_shape[2]) + split = resolve_ragged_pad_split( + bucket=int(bucket), + num_real_requests=n_real, + total_real_tokens=floor_tokens, + padded_bs=int(padded_bs), + max_verify_len=max_verify_len, + fixed_pad_len=fixed_pad_len, + ) + if split is None: + logger.warning( + f"DSpark ragged: bucket {bucket} admits no pad-row window " + f"for {n_real} real requests and {n_pad} pad rows; falling " + f"back to uniform scheduling" + ) + return None + pad_len = split.pad_len + real_target = split.real_target + if real_target != sum(token_lens): + logger.warning( + "DSpark exact layout changed between policy and fit: " + f"real tokens={sum(token_lens)}, target={real_target}, " + f"G={padded_bs}, V={bucket}, pad={pad_len}" + ) + return None + # The exact selector already returned a bounded split whose real rows + # sum to the measured cell; repacking it cannot change a value. + published = token_lens + + # Published only now, past the last way this fit can fail: a stale + # published bucket on a fallback step is the state behind the ragged + # IMA. The row count and pad window feed the fresh-confidence device + # prologue; `ragged_pad_verify_len` is also stamped on the shared dummy + # by _get_padded_batch. + self._dspark_last_padded_bs = int(padded_bs) + self._dspark_last_num_real = int(n_real) + runner.agreed_ragged_bucket = int(bucket) + runner.ragged_pad_verify_len = pad_len - 1 + + for request, tokens in zip(generation_requests, published): + request.py_verify_len = int(tokens) - 1 + return int(bucket) + + def _get_spec_worker(self): + """Return the one-model speculative worker when it is initialized.""" + return getattr(self.model, "spec_worker", None) + + def _apply_device_window_prologue(self, inputs, new_tensors_device) -> bool: + """Re-rank this step's verify windows on device, with fresh confidence. + + Runs after ``_prepare_inputs`` and before the graph replay, all device + ops on the current stream -- stream order makes every write visible to + the replayed graph, and nothing here reads device data back to the + host. The host has already agreed ``(padded_bs, bucket)`` and staged a + SHAPE SPLIT of the (lagged) budget through the normal fit; this + prologue re-distributes the same real/pad token totals by the verified + block's OWN confidence (ranking lag zero) and overwrites the layout's + content: verify lens, qo_indptr, the per-token row maps, the packed + input/position/draft tokens, and the per-request kv-length delta from + the shape split to the true windows. + + Returns True when applied; False when a precondition fails (the step + then runs the shape split as-is, which is a valid window assignment). + """ + from tensorrt_llm._torch.speculative.dspark_select import ( + gather_packed_draft_tokens, + select_windows_device, + ) + + runner = self.cuda_graph_runner + budget = self._dspark_device_budget + self._dspark_device_budget = None + if budget is None or runner.agreed_ragged_bucket is None: + return False + if not getattr(self, "_dspark_prev_covers_batch", False): + # The input gathers below address new_tokens_device by + # previous_batch_indices in batch order; a request without a + # previous device tensor breaks that addressing. + return False + worker = self._get_spec_worker() + if worker is None: + return False + planner = worker.verify_planner + if planner is None or worker.staged_confidence_buffer() is None: + return False + if worker.batch_slot_view(1) is None: + return False + spec_metadata = inputs.get("spec_metadata") + attn_metadata = inputs.get("attn_metadata") + if spec_metadata is None or attn_metadata is None: + return False + if self._config.use_mrope: + # Host-selected ragged windows build the complete three-axis + # MRoPE positions correctly. The device prologue currently owns + # only the scalar position buffer, so decline fresh reranking + # rather than leave the 3-D positions on the stale shape split. + return False + apply_device_layout = getattr(attn_metadata, "apply_device_ragged_layout", None) + if not callable(apply_device_layout): + # Device-selected row ownership is an attention-backend + # capability. Decline before mutating any shared buffers when a + # backend only supports the ordinary host-selected ragged layout. + return False + if int(attn_metadata.num_contexts) != 0: + return False + + bucket = int(runner.agreed_ragged_bucket) + padded_bs = int(self._dspark_last_padded_bs) + n_real = int(self._dspark_last_num_real) + pad_len_tok = int(runner.ragged_pad_verify_len) + 1 + real_tokens = bucket - (padded_bs - n_real) * pad_len_tok + cfg = planner.cfg + # The scheduler cannot grant more than the real rows can absorb under + # the published split; the fill tops up any shortfall. + budget = max(0, min(int(budget), real_tokens - n_real * (cfg.min_verify_len + 1))) + + # Snapshot the shape split BEFORE overwriting: past_seen per row is + # the staged position at each row's first token, and the kv delta + # needs the split the host baked into kv_lens_cuda. + lens_buf = self._buffers.ragged_verify_lens_cuda + qo_buf = self._buffers.ragged_qo_indptr_cuda + split_lens = lens_buf[:padded_bs].clone() + split_qo = qo_buf[: padded_bs + 1].to(torch.long) + past_seen = self._buffers.position_ids_cuda[split_qo[:-1]].clone() + + expected_stamp = worker.verified_draft_seq_cuda() + stamps = worker.confidence_stamp_buffer() if expected_stamp is not None else None + result = select_windows_device( + confidence_logits=worker.staged_confidence_buffer(), + slot_idx=worker.batch_slot_view(padded_bs), + num_real=n_real, + budget=budget, + graph_num_tokens=bucket, + cfg=cfg, + apply_calibration=planner.apply_calibration, + stamp=stamps, + expected_stamp=expected_stamp, + pad_len=pad_len_tok, + ) + + lens_buf[:padded_bs].copy_(result.verify_lens) + qo_buf[: padded_bs + 1].copy_(result.qo_indptr) + # The host stages kv_lens as num_cached + seq_lens_kv, and BOTH terms + # bake the per-request token window (num_cached = past + tokens; + # seq_lens_kv = tokens), so kv_lens = past + 2*S -- the window counts + # TWICE. Moving to the true windows therefore needs 2*(w - S), not + # (w - S): the single-delta variant left every re-ranked request's + # kv_len off by (w - S), which shifted the indexer K-cache slot + # mapping (slot_mapping_fp8) and silently wrote K entries into the + # wrong cache slots. Established empirically by an A/B tensor diff + # against a full host restage with the same windows. + window_delta = result.verify_lens - split_lens + attn_metadata.kv_lens_cuda[:padded_bs] += 2 * window_delta + # The graph adds previous_kv_lens_offsets (staged as new_tokens_lens - + # shape_lens, per request) to kv_lens during replay; host-with-w + # stages new_tokens_lens - w there, so the offsets move by -(w - S). + # Combined: (past + 2S) + 2(w-S) + (new - S) - (w-S) = past + w + new, + # exactly the host-with-w in-graph sum. + self._buffers.previous_kv_lens_offsets_cuda[:padded_bs] -= window_delta.to( + self._buffers.previous_kv_lens_offsets_cuda.dtype + ) + + req_idx = result.req_idx + spec_metadata.remap_expanded_sampling_params(req_idx, bucket) + device = req_idx.device + flat = torch.arange(bucket, device=device) + offset = flat - result.qo_indptr.to(torch.long)[req_idx] + prev_slots = self._buffers.previous_batch_indices_cuda[:n_real].to(torch.long) + slots_tok = prev_slots[req_idx.clamp(max=n_real - 1)] + + new_tokens_device = new_tensors_device.new_tokens + new_tokens_lens_device = new_tensors_device.new_tokens_lens + next_draft_tokens_device = new_tensors_device.next_draft_tokens + + self._buffers.input_ids_cuda[:bucket] = ( + new_tokens_device.transpose(0, 1)[slots_tok, offset] + .flatten() + .to(self._buffers.input_ids_cuda.dtype) + ) + self._buffers.position_ids_cuda[:bucket] = past_seen[req_idx] + offset.to(past_seen.dtype) + # Overlap corrections gather by each token's OWNER; rebuild both the + # index and the per-token offset it feeds (mirrors the host staging at + # the previous_pos_indices block). + self._buffers.previous_pos_indices_cuda[:real_tokens] = slots_tok[:real_tokens].to( + self._buffers.previous_pos_indices_cuda.dtype + ) + self._buffers.previous_pos_id_offsets_cuda[:real_tokens].copy_( + new_tokens_lens_device[slots_tok[:real_tokens]] + ) + # Draft tokens pack compactly, omitting each request's bonus/anchor. + # Build the draft-only row owners at the statically known real-draft + # size. The previous scheme parked anchors at ``real_draft``; at a + # full K / full batch that index is exactly one past the allocation. + real_draft = real_tokens - n_real + if real_draft > 0: + self._buffers.draft_tokens_cuda[:real_draft].copy_( + gather_packed_draft_tokens( + next_draft_tokens=next_draft_tokens_device, + batch_slots=prev_slots, + verify_lens=result.verify_lens, + qo_indptr=result.qo_indptr, + num_real=n_real, + total_draft_tokens=real_draft, + ).to(self._buffers.draft_tokens_cuda.dtype) + ) + + apply_device_layout(result.verify_lens, req_idx, result.kv_correction) + return True + + @staticmethod + def _ragged_token_lens(generation_requests) -> list[int] | None: + """Each generation request's token window, or None if the batch is uniform. + + ``py_verify_len`` counts drafted positions; the token window adds the + bonus position. None unless *every* request carries a window -- + partially windowed batches are silent token misattribution. + """ + verify_lens = [getattr(request, "py_verify_len", None) for request in generation_requests] + if not verify_lens or any(v is None for v in verify_lens): + return None + return [1 + int(v) for v in verify_lens] + + def _publish_gen_token_layout( + self, attn_metadata, generation_requests, *, runtime_draft_len: int + ) -> None: + """Hand the attention metadata this step's gen-token layout, before it + prepares. + + Split out of :meth:`_attach_ragged_verify_layout` for ordering: + ``attn_metadata.prepare()`` is the only consumer of + ``ragged_verify_lens`` and runs well before the spec metadata is + assembled. The uniform stride is published here too, because + ``max_draft_tokens`` is a static buffer-sizing ceiling that does not + move when a shorter tier is chosen. It cannot live in + ``update_spec_dec_param`` either: that runs against the base metadata, + and ``prepare()`` runs against the per-key CUDA-graph copy. + """ + if attn_metadata is None: + return + if hasattr(attn_metadata, "runtime_tokens_per_gen_step"): + attn_metadata.runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step( + runtime_draft_len + ) + if hasattr(attn_metadata, "ragged_verify_lens"): + attn_metadata.ragged_verify_lens = ( + self._ragged_token_lens(generation_requests) + if self._dspark_trims_submitted_tokens + else None + ) + if hasattr(attn_metadata, "device_windows_mode"): + # Tells the attention prepare that the host window VALUES are a + # shape split (bounds only); the true windows land on device + # through apply_device_ragged_layout after prepare. + attn_metadata.device_windows_mode = ( + self._dspark_trims_submitted_tokens and self._dspark_device_windows + ) + + def _pinned_host(self, key: str, values, dtype) -> torch.Tensor: + """A persistent pinned staging buffer holding ``values``. + + Async H2D sources must outlive the queued copy -- PyTorch does not + extend the lifetime of an async source, so a temporary tensor can be + reclaimed while the DMA still reads it. Buffers are keyed by name and + grown monotonically. Two slots support the token and draft gathers + issued under the same key in one preparation pass. The executor calls + :meth:`wait_for_input_copy` before the next pass mutates host input, so + no additional per-slot event or synchronization is needed. + """ + values = list(values) + slot = 1 - self._pinned_host_active.get(key, 1) + bufs = self._pinned_host_cache.setdefault(key, [None, None]) + buf = bufs[slot] + if buf is None or buf.numel() < len(values) or buf.dtype != dtype: + buf = torch.empty(max(len(values), 1), dtype=dtype, pin_memory=prefer_pinned()) + bufs[slot] = buf + self._pinned_host_active[key] = slot + view = buf[: len(values)] + if values: + view.copy_(torch.tensor(values, dtype=dtype)) + return view + + def _attach_ragged_verify_layout( + self, spec_metadata, attn_metadata, generation_requests + ) -> None: + """Publish this step's per-request verify windows to both metadatas. + + Only the DSpark confidence scheduler sets ``py_verify_len``; every + other path leaves the fields None. Acceptance needs the windows to + slice correctly; the DSA indexer needs them to expand kv_lens/block + tables per query token. + """ + if not self._dspark_trims_submitted_tokens: + spec_metadata.verify_lens = None + spec_metadata.qo_indptr = None + spec_metadata.total_verify_tokens = None + if attn_metadata is not None and hasattr(attn_metadata, "ragged_verify_lens"): + attn_metadata.ragged_verify_lens = None + return + + token_lens = self._ragged_token_lens(generation_requests) + if token_lens is None: + spec_metadata.verify_lens = None + spec_metadata.qo_indptr = None + spec_metadata.total_verify_tokens = None + if attn_metadata is not None and hasattr(attn_metadata, "ragged_verify_lens"): + attn_metadata.ragged_verify_lens = None + return + + from tensorrt_llm._torch.speculative.ragged_helpers import build_qo_indptr + + n = len(token_lens) + # Persistent buffers: a captured graph baked in the address it saw at + # capture time, so a fresh tensor would be invisible to every replay. + lens_view = self._buffers.ragged_verify_lens_cuda[:n] + lens_view.copy_( + self._pinned_host("ragged_verify_lens", token_lens, torch.int32), non_blocking=True + ) + indptr_view = self._buffers.ragged_qo_indptr_cuda[: n + 1] + indptr_view.copy_(build_qo_indptr(lens_view), non_blocking=True) + spec_metadata.verify_lens = lens_view + spec_metadata.qo_indptr = indptr_view + spec_metadata.total_verify_tokens = sum(token_lens) + if attn_metadata is not None and hasattr(attn_metadata, "ragged_verify_lens"): + attn_metadata.ragged_verify_lens = token_lens + + def _ragged_gather_indices( + self, slots: list[int], counts: list[int] + ) -> tuple[torch.Tensor, torch.Tensor]: + """Row/column index pair for gathering a ragged block out of a + [num_slots, max_width] device tensor. + + ``slots[i]`` contributes ``counts[i]`` entries taken from columns + ``0..counts[i]-1``, concatenated in batch order. This replaces the + ``tensor[slots, :width]`` strided gather used when every request + verifies the same number of positions. + """ + rows, cols = ragged_gather_index_lists(slots, counts) + rows_dev = self._pinned_host("gather_rows", rows, torch.long).to("cuda", non_blocking=True) + cols_dev = self._pinned_host("gather_cols", cols, torch.long).to("cuda", non_blocking=True) + return (rows_dev, cols_dev) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py index 11797542225c..266d51671be1 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py @@ -19,6 +19,7 @@ ResourceManagerType, ) from tensorrt_llm._torch.speculative import SpecMetadata, get_spec_metadata +from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from ...metadata import build_attention_metadata @@ -96,6 +97,7 @@ def setup_attention( else None, kv_cache_manager=kv_cache_manager, draft_kv_cache_manager=draft_kv_cache_manager, + enable_ragged_verification=self._ragged_verification_enabled, ) if isinstance(kv_cache_manager, BaseMambaCacheManager): self.attn_metadata.mamba_chunk_size = getattr( @@ -115,6 +117,142 @@ def setup_speculative( ) return self.spec_metadata + @property + def _ragged_verification_enabled(self) -> bool: + spec_config = self._config.spec_config + return bool( + spec_config is not None and getattr(spec_config, "enable_confidence_scheduling", False) + ) + + @torch.inference_mode() + def _warmup_dspark_ragged_compressor_metadata(self) -> None: + """Compile DSv4 ragged metadata for a non-CUDA-graph batch size. + + Whole-model CUDA-graph warmup exercises only configured graph batch + sizes. The first mixed attention-DP iteration can contain an + intermediate number of generation requests, which otherwise triggers + a long ``torch.compile`` specialization immediately before a TP + collective. Ranks without generation work reach the collective first + and can trip the executor hang detector while their peers compile. + + One non-bucket probe forces the dynamic specialization during engine + warmup, before the hang detector starts. It uses the metadata's real + persistent buffers so the tensor shape and stride guards match serving. + Those buffers are created under inference mode, so all probe and + cleanup mutations must run under the same mode. + """ + if not self._ragged_verification_enabled: + return + + metadata = getattr(self, "attn_metadata", None) + required_attributes = ( + "_compute_compressed_mask", + "_compute_gen_compressed_position_ids", + "_compress_ratios_sorted", + "compressed_mask_cuda", + "compressed_position_ids_cuda", + "cu_new_comp_kv_cuda", + "batch_size_cuda", + "max_draft_tokens", + "new_comp_kv_lens_cuda", + "past_kv_lens_cuda", + ) + missing_attributes = [ + name for name in required_attributes if metadata is None or not hasattr(metadata, name) + ] + if missing_attributes: + raise RuntimeError( + "DSpark ragged metadata warmup requires the DSv4 compressor " + "contract; missing " + ", ".join(missing_attributes) + ) + + graph_batch_sizes = {int(batch_size) for batch_size in self._config.cuda_graph_batch_sizes} + num_generations = next( + ( + candidate + for candidate in range(int(self._config.max_batch_size) - 1, 0, -1) + if candidate not in graph_batch_sizes + ), + None, + ) + if num_generations is None: + return + + compress_ratios = [int(ratio) for ratio in metadata._compress_ratios_sorted] + tokens_per_generation = 1 + int(metadata.max_draft_tokens) + total_compressed_tokens = {} + gen_output_offsets = {ratio: 0 for ratio in compress_ratios} + prepared_ratios = [] + original_batch_size = metadata.batch_size_cuda.clone() + try: + for ratio in compress_ratios: + compressed_tokens_per_generation = (tokens_per_generation + ratio - 1) // ratio + total_tokens = num_generations * compressed_tokens_per_generation + past_kv_lens = metadata.past_kv_lens_cuda[ratio] + cu_new_comp = metadata.cu_new_comp_kv_cuda[ratio] + new_comp = metadata.new_comp_kv_lens_cuda[ratio] + compressed_positions = metadata.compressed_position_ids_cuda[ratio] + compressed_mask = metadata.compressed_mask_cuda[ratio] + if ( + past_kv_lens.numel() < num_generations + or cu_new_comp.numel() < num_generations + 1 + or new_comp.numel() < num_generations + or compressed_positions.numel() < total_tokens + or compressed_mask.numel() < total_tokens + ): + raise RuntimeError( + "DSpark ragged compressor metadata buffers cannot " + f"hold the compile probe with {num_generations} " + "generation requests and compression ratio " + f"{ratio}" + ) + + total_compressed_tokens[ratio] = total_tokens + prepared_ratios.append(ratio) + past_kv_lens[:num_generations].zero_() + new_comp[:num_generations].fill_(compressed_tokens_per_generation) + cu_new_comp[: num_generations + 1].copy_( + torch.arange( + num_generations + 1, dtype=cu_new_comp.dtype, device=cu_new_comp.device + ) + ) + cu_new_comp[: num_generations + 1].mul_(compressed_tokens_per_generation) + + logger.info( + "DSpark ragged metadata warmup: compiling one non-graph " + f"shape with {num_generations} generation requests" + ) + metadata._compute_gen_compressed_position_ids( + metadata.past_kv_lens_cuda, + metadata.cu_new_comp_kv_cuda, + metadata.compressed_position_ids_cuda, + 0, + num_generations, + tokens_per_generation, + compress_ratios, + gen_output_offsets, + ) + metadata.batch_size_cuda.fill_(num_generations) + metadata._compute_compressed_mask( + metadata.new_comp_kv_lens_cuda, + metadata.cu_new_comp_kv_cuda, + metadata.compressed_mask_cuda, + metadata.batch_size_cuda, + compress_ratios, + ) + torch.cuda.synchronize() + finally: + metadata.batch_size_cuda.copy_(original_batch_size) + for ratio in prepared_ratios: + total_tokens = total_compressed_tokens[ratio] + metadata.past_kv_lens_cuda[ratio][:num_generations].zero_() + metadata.new_comp_kv_lens_cuda[ratio][:num_generations].zero_() + metadata.cu_new_comp_kv_cuda[ratio][: num_generations + 1].zero_() + metadata.compressed_position_ids_cuda[ratio][:total_tokens].zero_() + metadata.compressed_mask_cuda[ratio][:total_tokens].zero_() + torch.cuda.synchronize() + logger.info("DSpark ragged metadata warmup complete") + def get_spec_managers( resource_manager: ResourceManager, diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 952f6fd51375..391fc4d50386 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -30,6 +30,7 @@ from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( CUDAGraphRunner, CUDAGraphRunnerConfig, + DSparkCUDAGraphRunner, KeyType, ) from tensorrt_llm._torch.pyexecutor.guided_decoder import CapturableGuidedDecoder @@ -197,6 +198,7 @@ def __init__( self._init_decoder_state() self._inputs = self._make_input_preparer(mm_item_scheduler) self._graphs = self._make_cuda_graphs() + self._inputs.cuda_graph_runner = self.cuda_graph_runner self._initialize_moe_workspace_reclaimer() self._warmup = DecoderWarmup( config=self._config, @@ -279,6 +281,7 @@ def _cuda_graphs_kwargs(self) -> dict[str, Any]: torch_compile_bypass=self._bypass_torch_compile, context_promotion_supported=self._context_graph_promotion_supported, warmup_inputs=self._buffers.input_ids_cuda, + inputs=self._inputs, ) @property @@ -395,7 +398,8 @@ def _init_decoder_state(self) -> None: self._dynamic_draft_len_mapping = compute_dynamic_draft_len_mapping(self._config) def _initialize_cuda_graph_runner(self) -> CUDAGraphRunner | None: - cuda_graph_runner = CUDAGraphRunner(self._cuda_graph_runner_config()) + runner_type = DSparkCUDAGraphRunner if self._dspark_confidence_enabled else CUDAGraphRunner + cuda_graph_runner = runner_type(self._cuda_graph_runner_config()) register_hooks = getattr(self._model, "register_cuda_graph_pre_replay_hooks", None) if register_hooks is not None: register_hooks(cuda_graph_runner) @@ -648,6 +652,13 @@ def _forward_decoder( self._inputs.record_promoted_iter_states(scheduled_requests) self._inputs.record_input_copy() + if ( + can_run_graph + and not is_dummy + and getattr(self._inputs, "_dspark_device_budget", None) is not None + ): + self._inputs._apply_device_window_prologue(inputs, new_tensors_device) + def run_model(model_inputs: dict[str, Any]) -> Any: with MoeLoadBalancerIterContext(moe_load_balancer): return self._forward_step( @@ -882,6 +893,16 @@ def _execute_logit_post_processors( logits_processors = getattr(request, "py_logits_post_processors", None) if logits_processors: + if ( + not is_context_request + and getattr(request, "py_verify_len", None) is not None + ): + raise RuntimeError( + "per-request logits post-processors are not supported with DSpark " + "ragged verification: the processor interface assumes one uniform " + "row stride and would modify another request's packed logits after " + "a short verify window" + ) token_ids = ( [request.get_tokens(0)] if is_context_request @@ -944,3 +965,55 @@ def _restore_attention_dp_semantic_outputs( # The sampler sees the original empty batch, never the execution row. # Copy the map because CUDA graph outputs are retained across replays. return {**outputs, "logits": outputs["logits"][:0]} + + # Scheduler-facing state stays on the input owner that stages each layout. + + @property + def _dspark_confidence_enabled(self): + return self._inputs._dspark_confidence_enabled + + @property + def _dspark_trims_submitted_tokens(self): + return self._inputs._dspark_trims_submitted_tokens + + @property + def _dspark_sps_cost_table(self): + return self._inputs._dspark_sps_cost_table + + @property + def _dspark_exact_candidate_cells(self): + return self._inputs._dspark_exact_candidate_cells + + @property + def _dspark_exact_identity_words(self): + return self._inputs._dspark_exact_identity_words + + @property + def _dspark_device_budget(self): + return getattr(self._inputs, "_dspark_device_budget", None) + + @_dspark_device_budget.setter + def _dspark_device_budget(self, value) -> None: + self._inputs._dspark_device_budget = value + + @property + def previous_request_ids(self): + return self._inputs.previous_request_ids + + @property + def previous_verify_lens(self): + return self._inputs.previous_verify_lens + + @property + def has_previous_device_draft(self): + return self._inputs.has_previous_device_draft + + def ragged_verify_token_buckets(self, padded_bs: int): + return self._inputs.ragged_verify_token_buckets(padded_bs) + + def fit_ragged_verify_lens(self, *args, **kwargs): + return self._inputs.fit_ragged_verify_lens(*args, **kwargs) + + def _get_spec_worker(self): + """Return the one-model speculative worker when it is initialized.""" + return getattr(self._model, "spec_worker", None) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py index b9146c9c6bc3..4b5f99651907 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py @@ -481,6 +481,8 @@ def _warmup_scheduled( # creates; build it when every forward above was skipped. with self._warmup_timer.phase("dsa_prewarm"): self._ensure_dsa_attn_metadata_for_warmup(resource_manager) + if getattr(self._inputs, "_dspark_confidence_enabled", False): + self._metadata._warmup_dspark_ragged_compressor_metadata() with timing_metric("dg_paged_mqa_warmup_seconds", self._metrics): self._warmup_dg_paged_mqa_logits_metadata() log_mem_snapshot("warmup/after_dg_paged_mqa_logits_metadata") diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b76bd860e968..956e5c52b4f1 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -803,6 +803,61 @@ def cuda_graph_runner(self) -> Optional[CUDAGraphRunner]: return None return self._runner.cuda_graph_runner + @property + def _dspark_confidence_enabled(self) -> bool: + return isinstance( + self._runner, + DecoderRunner) and self._runner._dspark_confidence_enabled + + @property + def _dspark_trims_submitted_tokens(self) -> bool: + return isinstance( + self._runner, + DecoderRunner) and self._runner._dspark_trims_submitted_tokens + + @property + def _dspark_sps_cost_table(self): + if not isinstance(self._runner, DecoderRunner): + return None + return self._runner._dspark_sps_cost_table + + @property + def _dspark_exact_candidate_cells(self): + if not isinstance(self._runner, DecoderRunner): + return () + return self._runner._dspark_exact_candidate_cells + + @property + def _dspark_exact_identity_words(self): + if not isinstance(self._runner, DecoderRunner): + return (0, ) * 8 + return self._runner._dspark_exact_identity_words + + @property + def _dspark_device_budget(self): + if not isinstance(self._runner, DecoderRunner): + return None + return getattr(self._runner, "_dspark_device_budget", None) + + @_dspark_device_budget.setter + def _dspark_device_budget(self, value) -> None: + if isinstance(self._runner, DecoderRunner): + self._runner._dspark_device_budget = value + + def ragged_verify_token_buckets(self, batch_size: int): + if not isinstance(self._runner, DecoderRunner): + return () + return self._runner.ragged_verify_token_buckets(batch_size) + + def fit_ragged_verify_lens(self, *args, **kwargs): + assert isinstance(self._runner, DecoderRunner) + return self._runner.fit_ragged_verify_lens(*args, **kwargs) + + def _get_spec_worker(self): + if isinstance(self._runner, DecoderRunner): + return self._runner._get_spec_worker() + return None + @property def metrics(self) -> dict[str, float]: """Return model-engine warmup time metrics.""" diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 1839de961f8e..92d69414414e 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -246,8 +246,8 @@ def prepare(self): # without this, all concurrent gen requests fall through to the shared # scratch row below and corrupt each other's draft window at batch # size > 1 (GitHub #16767). Context-prefix entries are left to - # ``_seed_context_windows``. CUDA-graph capture dummies and real - # generation requests acquire a persistent slot only once; prepare() + # ``_seed_context_windows``. Real generation requests acquire a + # persistent slot only once; prepare() # runs before every replay, so it must never reset a live slot. # ADP-idle (id 0) and high-ID CUDA-graph padding dummies use scratch. num_contexts = max(0, len(self.request_ids) - self.num_generations) @@ -559,8 +559,9 @@ def _lazy_init(self, draft_model, spec_metadata, attn_metadata=None) -> None: self._scratch_slot = max_batch num_rows = max_batch + 1 - # CUDA-graph padding requests carry ids in - # ``[CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len, CUDA_GRAPH_DUMMY_REQUEST_ID]``, + # Both CUDA-graph padding variants, each with K+1 runtime widths, + # use ``[CUDA_GRAPH_DUMMY_REQUEST_ID - 2*(K+1) + 1, + # CUDA_GRAPH_DUMMY_REQUEST_ID]``, where K is the physical draft cap, # while real request ids start at ``max_batch_size`` and grow, so a simple # floor cleanly separates them. Together with ``ATTENTION_DP_DUMMY_REQUEST_ID`` # (0) these dummies must route to the scratch row (see ``prepare()``) and @@ -568,7 +569,9 @@ def _lazy_init(self, draft_model, spec_metadata, attn_metadata=None) -> None: # dspark -> cuda_graph_runner -> speculative.utils -> dspark import cycle. from ..pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID - self._graph_dummy_id_floor = CUDA_GRAPH_DUMMY_REQUEST_ID - self.max_draft_len + self._graph_dummy_id_floor = ( + CUDA_GRAPH_DUMMY_REQUEST_ID - 2 * (self.max_draft_len + 1) + 1 + ) self._kv_windows = torch.zeros( (num_rows, num_stages, self._win, head_dim), diff --git a/tensorrt_llm/_torch/speculative/dspark_select.py b/tensorrt_llm/_torch/speculative/dspark_select.py new file mode 100644 index 000000000000..800a6e81ac8d --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_select.py @@ -0,0 +1,206 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""DSpark verify-window selection: the in-graph prologue. + +The host planner decides windows from a confidence snapshot that is one +iteration old (the batch reshuffles across the lag, and the read itself must +not sync). This module ranks the block that is about to be verified by its +OWN confidence instead: everything from the slot-indexed gather to the +per-token row maps is a pure tensor function, so it can run at the head of +the captured graph, after the previous step's draft wrote the confidence +buffer and before this step's verify attention reads the layout. + +The split of responsibilities follows the paper's dual-timescale design: + +* the CAPACITY -- ``(padded_bs, bucket, budget)`` -- stays host-decided from + the lagged snapshot, because it is the CUDA-graph key and the attention-DP + agreement payload, both of which must exist before launch; +* the RANKING -- which requests win the budget -- is computed here from + fresh confidence, at zero staleness. + +Nothing in this module raises on data: a captured graph cannot branch. +Feasibility (the bucket bounds) is checkable from host-known scalars before +launch; per-row staleness degrades to the neutral row (verify everything for +that request), mirroring the host planner's fail-open semantics. +""" + +from dataclasses import dataclass +from typing import Callable, Optional + +import torch + +from .dspark_schedule import ( + NEUTRAL_CONFIDENCE_LOGIT, + DSparkScheduleConfig, + compute_survival, + schedule_verify_lens_topk, +) +from .ragged_helpers import build_qo_indptr, build_row_maps_device, fill_bucket_device + +__all__ = [ + "DeviceWindowResult", + "gather_packed_draft_tokens", + "select_windows_device", +] + + +@dataclass +class DeviceWindowResult: + """Everything the ragged layout consumers need, all device-resident. + + Attributes: + verify_lens: ``[padded_bs]`` int32 token windows (bonus included), + summing to exactly ``graph_num_tokens``; pad rows carry their + fill. + qo_indptr: ``[padded_bs + 1]`` int32 exclusive prefix sum. + req_idx: ``[graph_num_tokens]`` int64 owning-row per packed token. + kv_correction: ``[graph_num_tokens]`` int32; composed with a + ``kv_lens`` gather it yields each token's KV extent + (``refresh_ragged_row_kv_lens``). + """ + + verify_lens: torch.Tensor + qo_indptr: torch.Tensor + req_idx: torch.Tensor + kv_correction: torch.Tensor + + +def gather_packed_draft_tokens( + *, + next_draft_tokens: torch.Tensor, + batch_slots: torch.Tensor, + verify_lens: torch.Tensor, + qo_indptr: torch.Tensor, + num_real: int, + total_draft_tokens: int, +) -> torch.Tensor: + """Gather the device-selected real draft rows, excluding bonus tokens. + + ``verify_lens`` and ``qo_indptr`` describe token windows that include one + bonus/anchor per request. The persistent draft buffer contains drafts + only, packed request-major. Construct exactly ``total_draft_tokens`` + owners so full-batch/full-K layouts never need a one-past-the-end discard + slot for anchors. + """ + if total_draft_tokens < 0: + raise ValueError("total_draft_tokens must be non-negative") + if total_draft_tokens == 0: + return next_draft_tokens.new_empty((0,)) + device = verify_lens.device + rows = torch.arange(num_real, device=device, dtype=torch.long) + draft_counts = verify_lens[:num_real].to(torch.long) - 1 + owners = torch.repeat_interleave(rows, draft_counts, output_size=total_draft_tokens) + # Removing one bonus from every preceding request converts the token + # prefix into a draft-only prefix. + draft_qo = qo_indptr[: num_real + 1].to(torch.long) - torch.arange( + num_real + 1, device=device, dtype=torch.long + ) + flat = torch.arange(total_draft_tokens, device=device, dtype=torch.long) + offsets = flat - draft_qo[owners] + slots = batch_slots[:num_real].to(torch.long)[owners] + return next_draft_tokens[slots, offsets] + + +def select_windows_device( + *, + confidence_logits: torch.Tensor, + slot_idx: torch.Tensor, + num_real: torch.Tensor, + budget: torch.Tensor, + graph_num_tokens: int, + cfg: DSparkScheduleConfig, + apply_calibration: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, + stamp: Optional[torch.Tensor] = None, + expected_stamp: Optional[torch.Tensor] = None, + pad_len: Optional[int] = None, +) -> DeviceWindowResult: + """Rank the batch by fresh confidence and pack it into the agreed bucket. + + Mirrors the host chain ``_gather_rows -> apply_calibration -> + compute_survival -> schedule_verify_lens_topk -> fill_bucket`` with the + lag removed: ``confidence_logits`` is the LIVE slot-indexed buffer the + previous draft pass scattered into, not a staged snapshot. + + Args: + confidence_logits: ``[num_slots, K]`` raw logits, slot-indexed. + slot_idx: ``[padded_bs]`` int64 buffer row per batch position; + entries at or beyond ``num_real`` are never read meaningfully + (point them at any valid row, e.g. 0). + num_real: 0-d integer tensor; rows past it are padding. + budget: 0-d integer tensor, verify tokens above the floor -- the + host-decided (lagged) capacity knob, written before launch. + graph_num_tokens: the captured token bucket (capture constant). + cfg: scheduling bounds; ``resolved_max_verify_len + 1`` is the + per-row token ceiling. + apply_calibration: logits -> probabilities; sigmoid when None. + stamp: optional ``[num_slots]`` last-writer stamps; with + ``expected_stamp`` (0-d), rows whose stamp mismatches fall back + to the neutral row (full survival: verify everything), the same + fail-open the host planner applies to unknown slots. + pad_len: when given, pad rows carry EXACTLY this many tokens and all + fill slack goes to real rows -- the host fit's published split, + which its pre-launch copy widths depend on (see + :func:`fill_bucket_device`). The caller must clamp ``budget`` so + the real rows can absorb ``graph_num_tokens - n_pad * pad_len``. + + Returns: + :class:`DeviceWindowResult`; every tensor is derived without a host + sync, so the call is capture-safe end to end. + """ + if confidence_logits.dim() != 2: + raise ValueError( + f"confidence_logits must be [num_slots, K], got {tuple(confidence_logits.shape)}" + ) + device = confidence_logits.device + padded_bs = slot_idx.numel() + + selected = confidence_logits.index_select(0, slot_idx.to(torch.long)) + if stamp is not None: + if expected_stamp is None: + raise ValueError("stamp requires expected_stamp") + stale = stamp.index_select(0, slot_idx.to(torch.long)) != expected_stamp + selected = torch.where( + stale.unsqueeze(1), + torch.full_like(selected, NEUTRAL_CONFIDENCE_LOGIT), + selected, + ) + + calibrate = apply_calibration or torch.sigmoid + survival = compute_survival(calibrate(selected)) + # Pad rows must win no budget: zero survival sits below survival_eps, so + # the ranking never selects them and their windows stay at the floor + # until the fill tops them up. + is_real = torch.arange(padded_bs, device=device) < num_real + survival = torch.where(is_real.unsqueeze(1), survival, torch.zeros_like(survival)) + + scheduled = schedule_verify_lens_topk(survival=survival, budget=budget, cfg=cfg) + # The scheduler counts drafted positions; the token window adds the + # bonus/anchor, matching every host fill_bucket callsite. + token_lens = scheduled + 1 + max_token_len = int(cfg.resolved_max_verify_len) + 1 + filled = fill_bucket_device( + token_lens, + num_real=num_real, + graph_num_tokens=graph_num_tokens, + max_verify_len=max_token_len, + pad_fill=pad_len, + ) + req_idx, correction = build_row_maps_device(filled, graph_num_tokens=graph_num_tokens) + return DeviceWindowResult( + verify_lens=filled, + qo_indptr=build_qo_indptr(filled), + req_idx=req_idx, + kv_correction=correction, + ) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index eb9067c54f95..ce97ef6322eb 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -1574,6 +1574,52 @@ def invalidate_sampling_params_cache(self) -> None: self._sampling_params_signature[0] = None self._sampling_params_signature[1] = None + def remap_expanded_sampling_params(self, req_idx: torch.Tensor, + num_tokens: int) -> None: + """Repack per-logits-row sampling state after device window selection. + + Device window selection can redistribute a fixed token budget among + requests immediately before replay. The expanded sampling buffers must + follow the resulting token owners. + """ + if num_tokens <= 0: + return + owners = req_idx[:num_tokens].to(torch.long) + pairs = ( + (self.request_temperatures, self.temperatures), + (self.request_top_ks, self.top_ks), + (self.request_top_ps, self.top_ps), + (self.request_min_ps, self.min_ps), + (self.request_seeds, self.seeds), + (self.request_offsets, self.offsets), + ) + for request_values, expanded_values in pairs: + if request_values is None or expanded_values is None: + continue + if expanded_values.numel() < num_tokens: + raise RuntimeError( + "device-selected DSpark layout exceeds the captured " + "sampling-parameter buffer") + torch.index_select(request_values, + 0, + owners, + out=expanded_values[:num_tokens]) + + if self.request_offsets is not None and self.offsets is not None: + if self.qo_indptr is None: + raise RuntimeError( + "device-selected DSpark sampling requires ragged row starts" + ) + row_slots = torch.arange(num_tokens, + device=owners.device, + dtype=torch.int64) + row_slots -= self.qo_indptr.index_select(0, owners) + self.offsets[:num_tokens].add_(row_slots * _RNG_SLOT_SPAN) + + # The expanded cache now follows the device-selected layout. Force the + # next host-prepared iteration to restore its own row expansion. + self._sampling_params_signature[1] = None + class AuxiliarySpeculativeStateHandler(Protocol): """Model-side state that verification writes once per draft token. diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 3ed9f4085675..6d574e50bce2 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -115,6 +115,7 @@ l0_cpu: - unittest/_torch/multimodal - unittest/_torch/ray_orchestrator/single_gpu/test_cache_transceiver_comm.py - unittest/_torch/speculative/hw_agnostic + - unittest/_torch/speculative/test_device_window_sampling_layout.py # l0_h100 selects with -m "not cpu_only", so this entry is what runs it. - unittest/_torch/speculative/test_spec_slot_pool_sizing.py - unittest/_torch/test_custom_config_registration.py diff --git a/tests/unittest/_torch/compilation/test_phase_selective_forward.py b/tests/unittest/_torch/compilation/test_phase_selective_forward.py index 41ecf6738563..6ee9555bafb4 100644 --- a/tests/unittest/_torch/compilation/test_phase_selective_forward.py +++ b/tests/unittest/_torch/compilation/test_phase_selective_forward.py @@ -233,6 +233,7 @@ def test_decoder_runner_warmup_bypasses_only_auxiliary_phases(monkeypatch, enabl enabled=False, is_warmup_only=False, padding_dummy_requests={}, + reset_padding_dummies=Mock(), ) graphs.preallocate_padding_dummies = Mock() warmup = object.__new__(DecoderWarmup) @@ -242,6 +243,7 @@ def test_decoder_runner_warmup_bypasses_only_auxiliary_phases(monkeypatch, enabl warmup._warmup_timer = _WarmupTimer(rank=0) warmup._agreement = SimpleNamespace(agree_flag=lambda value: value, agree_shapes=lambda s: s) warmup._graphs = graphs + warmup._metadata = SimpleNamespace(_warmup_dspark_ragged_compressor_metadata=Mock()) warmup._context_warmups_supported = True warmup._torch_compile_bypass = runner._bypass_torch_compile warmup._get_full_general_warmup_requests = Mock(return_value=[(4, 0)]) diff --git a/tests/unittest/_torch/executor/test_moe_comm_timeout_warmup.py b/tests/unittest/_torch/executor/test_moe_comm_timeout_warmup.py index d56cefa658b0..f40d92aabaff 100644 --- a/tests/unittest/_torch/executor/test_moe_comm_timeout_warmup.py +++ b/tests/unittest/_torch/executor/test_moe_comm_timeout_warmup.py @@ -16,6 +16,7 @@ DEFAULT_WARMUP_TIMEOUT_SEC, MoECommTimeoutGuard, ) +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs import DecoderCudaGraphs from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine @@ -55,6 +56,8 @@ def __init__(self) -> None: class _GraphRunnerStub: + reset_padding_dummies = CUDAGraphRunner.reset_padding_dummies + def __init__(self) -> None: self.is_warmup_only = False self.padding_dummy_requests: dict[str, object] = {"stale": object()} diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index b26cb976be99..2fb36454174a 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -996,6 +996,7 @@ def test_graph_key_forwards_promoted_context_ids(self) -> None: runner.config = SimpleNamespace() runner._resolve_sample_type.return_value = SampleType.FULL runner._get_seq_len_mode.return_value = True + runner._ragged_verify_bucket.return_value = None request = _make_request_stub(7) batch = ScheduledRequests() batch.generation_requests = [request] @@ -1023,6 +1024,7 @@ def test_graph_key_aggregates_encoder_tokens(self) -> None: runner._resolve_sample_type.return_value = SampleType.FULL runner.max_beam_width = 1 runner._get_seq_len_mode.return_value = False + runner._ragged_verify_bucket.return_value = None context = _make_request_stub(1) context.encoder_output_len = 7 context.py_skip_cross_kv_projection = False @@ -1047,6 +1049,7 @@ def test_graph_key_rejects_nonuniform_context_query_lengths(self) -> None: runner.config = SimpleNamespace() runner._resolve_sample_type.return_value = SampleType.FULL runner._get_seq_len_mode.return_value = False + runner._ragged_verify_bucket.return_value = None first_context = _make_request_stub(1) first_context.encoder_output_len = 7 first_context.py_skip_cross_kv_projection = False @@ -1093,6 +1096,7 @@ def test_graph_key_includes_peft_cache_dtype(self) -> None: runner.config = SimpleNamespace() runner._resolve_sample_type.return_value = SampleType.FULL runner._get_seq_len_mode.return_value = False + runner._ragged_verify_bucket.return_value = None request = _make_request_stub(7) batch = ScheduledRequests() batch.generation_requests = [request] @@ -1111,6 +1115,9 @@ def test_graph_key_includes_peft_cache_dtype(self) -> None: def test_graph_dtype_change_falls_back_to_eager(self) -> None: runner = Mock() runner.enabled = True + runner._agree_graph_shape = CUDAGraphRunner._agree_graph_shape.__get__( + runner) + runner._dspark_confidence_enabled = False runner.config = SimpleNamespace( enable_attention_dp=False, use_mrope=False, @@ -1149,6 +1156,7 @@ def test_graph_key_includes_lora_variant(self) -> None: runner.config = SimpleNamespace() runner._resolve_sample_type.return_value = SampleType.FULL runner._get_seq_len_mode.return_value = False + runner._ragged_verify_bucket.return_value = None request = _make_request_stub(7) batch = ScheduledRequests() batch.generation_requests = [request] @@ -1188,6 +1196,9 @@ def test_lora_graph_variant_selection(self) -> None: def test_graph_lookup_forwards_promoted_context_ids(self) -> None: runner = Mock() runner.enabled = True + runner._agree_graph_shape = CUDAGraphRunner._agree_graph_shape.__get__( + runner) + runner._dspark_confidence_enabled = False runner.config = SimpleNamespace( enable_attention_dp=False, use_mrope=False, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_cuda_graph_dummy_identity.py b/tests/unittest/_torch/speculative/hw_agnostic/test_cuda_graph_dummy_identity.py new file mode 100644 index 000000000000..82ca8377d59d --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_cuda_graph_dummy_identity.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( + CUDA_GRAPH_DUMMY_REQUEST_ID, + _padding_dummy_draft_token_capacity, + cuda_graph_dummy_request_id, +) +from tensorrt_llm._torch.speculative import dspark as dspark_module +from tensorrt_llm._torch.speculative.dspark import DSparkSpecMetadata, DSv4DSparkWorker + + +@pytest.mark.parametrize("mode", ["Eagle", "MTP"]) +def test_tree_speculation_keeps_native_dummy_identity(mode): + config = SimpleNamespace(decoding_type=mode, max_draft_len=3) + capacity = _padding_dummy_draft_token_capacity(config, 12) + assert capacity == 12 + assert cuda_graph_dummy_request_id(12, variant=0, max_draft_len=capacity) == ( + CUDA_GRAPH_DUMMY_REQUEST_ID - 12 + ) + + +def test_dspark_does_not_relax_the_physical_k_bound(): + config = SimpleNamespace(decoding_type="DSpark", max_draft_len=3) + capacity = _padding_dummy_draft_token_capacity(config, 12) + assert capacity == 3 + with pytest.raises(ValueError, match="invalid CUDA padding dummy identity"): + cuda_graph_dummy_request_id(12, variant=0, max_draft_len=capacity) + + +def test_secondary_dspark_dummy_ids_are_disjoint(): + native = {cuda_graph_dummy_request_id(k, variant=0, max_draft_len=3) for k in range(4)} + secondary = {cuda_graph_dummy_request_id(k, variant=1, max_draft_len=3) for k in range(4)} + assert native.isdisjoint(secondary) + + +def test_target_only_dummy_identity_is_unchanged(): + capacity = _padding_dummy_draft_token_capacity(None, 0) + assert cuda_graph_dummy_request_id(0, variant=0, max_draft_len=capacity) == ( + CUDA_GRAPH_DUMMY_REQUEST_ID + ) + + +@pytest.mark.parametrize("physical_k", [1, 3, 8]) +def test_padding_families_never_consume_real_worker_slots(monkeypatch, physical_k): + def cpu_allocator(original): + def allocate(*args, **kwargs): + kwargs["device"] = "cpu" + return original(*args, **kwargs) + + return allocate + + for name in ("empty", "zeros", "full"): + monkeypatch.setattr(torch, name, cpu_allocator(getattr(torch, name))) + monkeypatch.setattr(dspark_module, "prefer_pinned", lambda: False) + worker = object.__new__(DSv4DSparkWorker) + worker.spec_config = SimpleNamespace(max_draft_len=physical_k) + worker._win_inited = False + worker._attention_warmup_attempted = True + worker.return_confidence = False + worker._confidence_logits = None + draft = SimpleNamespace( + block_size=physical_k, + num_stages=1, + _attn_params={"window_size": 4, "head_dim": 2}, + ) + worker._lazy_init(draft, SimpleNamespace(max_num_requests=4)) + metadata = SimpleNamespace( + _dspark_worker=worker, + batch_indices_cuda=torch.empty(4, dtype=torch.int32), + ) + + for variant in (0, 1): + for runtime_k in range(physical_k + 1): + dummy_id = cuda_graph_dummy_request_id( + runtime_k, variant=variant, max_draft_len=physical_k + ) + metadata.request_ids = [1000, 0, dummy_id] + metadata.num_generations = 3 + DSparkSpecMetadata.prepare(metadata) + assert worker._req_to_slot == {1000: 0} + assert list(worker._free_slots) == [1, 2, 3] + assert worker._batch_to_slot[:3].tolist() == [0, 4, 4] + + metadata.request_ids = [0, dummy_id] + metadata.num_generations = 2 + DSparkSpecMetadata.prepare(metadata) + assert worker._req_to_slot == {} + assert sorted(worker._free_slots) == [0, 1, 2, 3] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph_runtime.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph_runtime.py new file mode 100644 index 000000000000..fb5e02160175 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph_runtime.py @@ -0,0 +1,301 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU tests for DSpark ragged CUDA-graph keys and padding state.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor import cuda_graph_runner as graph_module +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( + CUDA_GRAPH_DUMMY_REQUEST_ID, + ADPShapeAgreement, + CUDAGraphRunner, + DSparkCUDAGraphRunner, + KeyType, + cuda_graph_dummy_request_id, +) +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.sampler.sampler_common import SampleType + + +class _Batch: + def __init__(self, generation_requests): + self.generation_requests = list(generation_requests) + self.context_requests = [] + self.encoder_requests = [] + self.can_run_cuda_graph = True + + @property + def batch_size(self): + return len(self.generation_requests) + + @property + def num_context_requests(self): + return 0 + + @property + def num_generation_requests(self): + return len(self.generation_requests) + + +def _runner_stub(runner_type=DSparkCUDAGraphRunner): + runner = runner_type.__new__(runner_type) + runner.clear = lambda: None + return runner + + +def _agreement(batch, *, iteration=9, peer_sizes=(3, 4), reuse=True): + return ADPShapeAgreement( + iteration=iteration, + batch_identity=id(batch), + local_batch_size=batch.batch_size, + peer_batch_sizes=peer_sizes, + all_can_graph=True, + widest_batch_size=max(peer_sizes), + graph_batch_size=4, + draft_len=5, + padding_ready=True, + ragged_bucket=16, + reuse_graph_shape=reuse, + ) + + +def test_padding_dummy_ids_are_disjoint(): + ids = { + cuda_graph_dummy_request_id(draft_len, variant=variant, max_draft_len=5) + for variant in (0, 1) + for draft_len in range(6) + } + + assert len(ids) == 12 + assert ( + cuda_graph_dummy_request_id(5, variant=0, max_draft_len=5) + == CUDA_GRAPH_DUMMY_REQUEST_ID - 5 + ) + with pytest.raises(ValueError, match="invalid CUDA padding dummy identity"): + cuda_graph_dummy_request_id(5, variant=2, max_draft_len=5) + + +def test_shape_agreement_is_bound_to_iteration_batch_and_phase(): + batch = _Batch([object()] * 3) + agreement = _agreement(batch) + + assert agreement.matches(batch, 9, padded=False) + assert not agreement.matches(batch, 8, padded=False) + assert not agreement.matches(_Batch([object()] * 3), 9, padded=False) + + batch.generation_requests.append(object()) + assert agreement.matches(batch, 9, padded=True) + assert not agreement.matches(batch, 9, padded=False) + + +def test_ragged_graph_key_requires_every_row_and_the_agreed_bucket(): + runner = _runner_stub() + runner.agreed_ragged_bucket = 22 + batch = _Batch( + [ + SimpleNamespace(py_verify_len=3), + SimpleNamespace(py_verify_len=5), + ] + ) + + assert runner._ragged_verify_bucket(batch) == 22 + batch.generation_requests[1].py_verify_len = None + assert runner._ragged_verify_bucket(batch) is None + + +def test_native_runner_does_not_inspect_ragged_request_windows(): + runner = _runner_stub(CUDAGraphRunner) + assert runner._ragged_verify_bucket(object()) is None + + +def test_ragged_graph_key_controls_the_submitted_token_count(): + runner = _runner_stub() + + assert ( + runner._get_num_tokens_for_key( + KeyType( + batch_size=4, + draft_len=5, + ragged_verify_bucket=14, + ) + ) + == 14 + ) + with pytest.raises(ValueError, match="mixed context"): + runner._get_num_tokens_for_key( + KeyType( + batch_size=4, + draft_len=5, + num_contexts=1, + ragged_verify_bucket=14, + ) + ) + + +@pytest.mark.parametrize( + ("padded_size", "num_requests", "expected"), + [(4, 3, True), (4, 4, True), (8, 3, False), (0, 3, False), (4, 0, False)], +) +def test_will_pad_to_matches_the_size_guards(padded_size, num_requests, expected): + runner = _runner_stub() + runner.supported_batch_sizes = [1, 2, 4] + runner.config = SimpleNamespace(batch_size=4) + + assert runner.will_pad_to(padded_size, num_requests) is expected + + +def test_zero_real_padding_uses_distinct_low_and_high_dummies(): + runner = _runner_stub() + runner.enabled = True + runner.padding_enabled = True + runner.max_supported_batch_size = 4 + runner.config = SimpleNamespace( + enable_attention_dp=False, + mapping=SimpleNamespace(tp_size=1), + batch_size=4, + ) + runner.adp_shape_agreement = None + runner.ragged_pad_verify_len = 2 + runner.ragged_zero_real_high_rows = 2 + runner.spec_config = SimpleNamespace(enable_confidence_scheduling=True, max_draft_len=5) + runner._can_run_cuda_graph_batch = lambda _batch: True + runner._round_up_batch_size_with_draft_len = lambda *_args: 4 + + scheduled_dummy = SimpleNamespace(is_attention_dp_dummy=True, py_verify_len=3) + low = SimpleNamespace(py_verify_len=None) + high = SimpleNamespace(py_verify_len=None) + runner._get_or_create_padding_dummy = ( + lambda _resource, _draft, variant=0: low if variant == 0 else high + ) + batch = _Batch([scheduled_dummy]) + + added = runner._get_padded_batch(batch, SimpleNamespace(), runtime_draft_len=5) + + assert added == 3 + assert batch.generation_requests == [scheduled_dummy, high, low, low] + assert [row.py_verify_len for row in batch.generation_requests] == [3, 3, 2, 2] + assert sum(1 + row.py_verify_len for row in batch.generation_requests) == 14 + + +def test_releasing_padding_dummies_invalidates_the_cached_agreement(): + runner = _runner_stub() + low = object() + high = object() + runner.padding_dummy_requests = {5: low} + runner.secondary_padding_dummy_requests = {5: high} + runner.adp_shape_agreement = object() + freed = [] + runner._padding_dummy_managers = lambda _resource_manager: [ + SimpleNamespace(free_resources=freed.append) + ] + + assert runner.release_padding_dummy(SimpleNamespace(), 5) + assert freed == [low, high] + assert runner.adp_shape_agreement is None + + +@pytest.mark.parametrize("confidence_enabled", [False, True]) +def test_decoder_selects_confidence_graph_subclass_only_when_enabled( + monkeypatch, confidence_enabled +): + monkeypatch.setattr(graph_module, "get_memory_buffers", lambda: None) + config = SimpleNamespace( + use_cuda_graph=False, + cuda_graph_padding_enabled=True, + cuda_graph_batch_sizes=[1, 4], + max_cuda_graph_batch_size=4, + max_beam_width=1, + spec_config=SimpleNamespace(enable_confidence_scheduling=confidence_enabled), + sparse_attention_config=None, + is_encoder_decoder=False, + enable_encoder_decoder_mixed_cuda_graph=False, + enable_in_graph_sampling=False, + cuda_graph_mem_pool=None, + dynamic_draft_len_mapping=None, + ) + registered = [] + decoder = SimpleNamespace( + _dspark_confidence_enabled=confidence_enabled, + _cuda_graph_runner_config=lambda: config, + model=SimpleNamespace(register_cuda_graph_pre_replay_hooks=registered.append), + ) + + runner = DecoderRunner._initialize_cuda_graph_runner(decoder) + runner.clear = lambda: None + + expected_type = DSparkCUDAGraphRunner if confidence_enabled else CUDAGraphRunner + assert type(runner) is expected_type + assert registered == [runner] + for name in ( + "agreed_ragged_bucket", + "adp_shape_agreement", + "ragged_pad_verify_len", + "ragged_zero_real_high_rows", + "secondary_padding_dummy_requests", + ): + assert hasattr(runner, name) is confidence_enabled + + +@pytest.mark.parametrize("runner_type", [CUDAGraphRunner, DSparkCUDAGraphRunner]) +def test_reset_padding_dummies_clears_only_runner_owned_state(runner_type): + runner = _runner_stub(runner_type) + runner.padding_dummy_requests = {5: object()} + if runner_type is DSparkCUDAGraphRunner: + runner.secondary_padding_dummy_requests = {5: object()} + runner.adp_shape_agreement = object() + runner.ragged_zero_real_high_rows = 2 + + runner.reset_padding_dummies() + + assert runner.padding_dummy_requests == {} + if runner_type is DSparkCUDAGraphRunner: + assert runner.secondary_padding_dummy_requests == {} + assert runner.adp_shape_agreement is None + assert runner.ragged_zero_real_high_rows == 0 + else: + assert not hasattr(runner, "secondary_padding_dummy_requests") + assert not hasattr(runner, "adp_shape_agreement") + + +def test_cached_shape_agreement_keeps_sampling_tier_collective(): + runner = _runner_stub() + batch = _Batch([SimpleNamespace(py_verify_len=3, py_draft_tokens=[0] * 5)] * 4) + runner.adp_shape_agreement = _agreement(batch, peer_sizes=(4, 4)) + runner.agreed_ragged_bucket = 16 + runner.enabled = True + runner.enable_in_graph_sampling = True + exchanges = [] + + def allgather(payload): + exchanges.append(payload) + peer = list(payload) + peer[-1] = SampleType.FULL.value + return torch.tensor([payload, peer]) + + runner.config = SimpleNamespace( + enable_attention_dp=True, + mapping=SimpleNamespace(tp_size=2), + dist=SimpleNamespace(tp_allgather_int64=allgather), + ) + runner._resolve_sample_type = lambda *_args: SampleType.FAST + + assert runner._agree_graph_shape(batch, True, frozenset()) == (True, SampleType.FULL) + assert exchanges == [[True, 4, 16, 5, SampleType.FAST.value]] + runner.enable_in_graph_sampling = False + exchanges.clear() + assert runner._agree_graph_shape(batch, True, frozenset()) == (True, None) + assert exchanges == [] + + +def test_cached_shape_agreement_rejects_a_changed_ragged_key(): + runner = _runner_stub() + batch = _Batch([SimpleNamespace(py_verify_len=3, py_draft_tokens=[0] * 5)] * 4) + runner.adp_shape_agreement = _agreement(batch, peer_sizes=(4, 4)) + runner.agreed_ragged_bucket = 15 + runner.enable_in_graph_sampling = False + + with pytest.raises(RuntimeError, match="no longer matches"): + runner._agree_graph_shape(batch, True, frozenset()) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py new file mode 100644 index 000000000000..54c258940dc1 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py @@ -0,0 +1,701 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Focused CPU tests for DSpark decoder-runner ragged runtime state.""" + +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch import utils as torch_utils +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.metadata import ( + DeepseekV4TrtllmAttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import DSAtrtllmAttentionMetadata +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner, DSparkCUDAGraphRunner +from tensorrt_llm._torch.pyexecutor.engine import metadata as metadata_module +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import graphs as graphs_module +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import inputs as decoder_module +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import ( + metadata as decoder_metadata_module, +) +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs import DecoderCudaGraphs +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import DecoderInputPreparer +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.metadata import DecoderMetadata +from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine +from tensorrt_llm._torch.speculative.ragged_helpers import ragged_gather_index_lists + + +def _runner(*, secondary=True): + return SimpleNamespace( + enabled=True, + agreed_ragged_bucket=None, + ragged_pad_verify_len=0, + ragged_zero_real_high_rows=0, + supported_batch_sizes=[4, 16], + secondary_padding_dummy_requests={5: object()} if secondary else {}, + _round_up_batch_size=lambda rows: 4 if rows <= 4 else 16, + will_pad_to=lambda *_args: True, + ) + + +def _engine(runner, buckets): + engine = object.__new__(DecoderInputPreparer) + engine.cuda_graph_runner = runner + engine._config = SimpleNamespace( + spec_config=SimpleNamespace(max_draft_len=5), + max_batch_size=8, + cuda_graph_batch_sizes=[1, 4, 8], + use_mrope=False, + ) + engine._dspark_confidence_enabled = True + engine._dspark_trims_submitted_tokens = True + engine._dspark_device_windows = False + engine._dspark_last_padded_bs = None + engine.ragged_verify_token_buckets = lambda _rows: list(buckets) + return engine + + +def test_pinned_host_ping_pongs_without_private_cuda_events() -> None: + engine = SimpleNamespace( + _pinned_host_cache={}, + _pinned_host_active={}, + ) + + first = DecoderInputPreparer._pinned_host(engine, "rows", [1, 2], torch.long) + second = DecoderInputPreparer._pinned_host(engine, "rows", [3, 4], torch.long) + third = DecoderInputPreparer._pinned_host(engine, "rows", [5, 6], torch.long) + + assert first.data_ptr() == third.data_ptr() + assert first.data_ptr() != second.data_ptr() + assert third.tolist() == [5, 6] + + +@pytest.mark.parametrize( + ("verifier_budget", "scheduled_window", "high_rows"), + [(80, 4, 0), (81, 5, 1), (88, 5, 8)], +) +def test_zero_real_exact_fit_uses_one_scheduled_dummy(verifier_budget, scheduled_window, high_rows): + runner = _runner() + engine = _engine(runner, [48, 64, 80, 81, 88, 96]) + dummy = SimpleNamespace(is_attention_dp_dummy=True, py_verify_len=None) + + bucket = DecoderInputPreparer.fit_ragged_verify_lens( + engine, + [dummy], + [scheduled_window], + peer_stats=[[16, 88, 1], [0, 0, 1]], + exact_shape=(16, verifier_budget, 5), + exact_zero_real=True, + ) + + assert bucket == verifier_budget + assert dummy.py_verify_len == scheduled_window + assert engine._dspark_last_num_real == 0 + assert runner.ragged_pad_verify_len == 4 + assert runner.ragged_zero_real_high_rows == high_rows + + +def test_zero_real_exact_fit_requires_the_secondary_dummy(): + runner = _runner(secondary=False) + engine = _engine(runner, [88]) + dummy = SimpleNamespace(is_attention_dp_dummy=True, py_verify_len=None) + + with pytest.raises(RuntimeError, match="disappeared after"): + DecoderInputPreparer.fit_ragged_verify_lens( + engine, + [dummy], + [5], + peer_stats=[[16, 88, 1], [0, 0, 1]], + exact_shape=(16, 88, 5), + exact_zero_real=True, + ) + + +def test_exact_fit_publishes_the_measured_bucket_and_pad_window(): + runner = _runner() + engine = _engine(runner, [8, 24]) + requests = [SimpleNamespace(py_verify_len=None) for _ in range(2)] + + bucket = DecoderInputPreparer.fit_ragged_verify_lens( + engine, + requests, + [2, 2], + peer_stats=[[2, 6, 1]], + exact_shape=(4, 8, 1), + ) + + assert bucket == 8 + assert [request.py_verify_len for request in requests] == [2, 2] + assert runner.agreed_ragged_bucket == 8 + assert runner.ragged_pad_verify_len == 0 + + +def test_full_k_bucket_preserves_the_native_static_graph(): + runner = _runner() + engine = _engine(runner, [24]) + requests = [SimpleNamespace(py_verify_len=None) for _ in range(2)] + + bucket = DecoderInputPreparer.fit_ragged_verify_lens( + engine, + requests, + [5, 5], + peer_stats=[[2, 12, 1]], + exact_shape=(4, 24, 6), + ) + + assert bucket is None + assert runner.agreed_ragged_bucket is None + assert all(request.py_verify_len is None for request in requests) + + +def _warmup_metadata(position_helper, mask_helper): + ratios = [1, 4] + return SimpleNamespace( + _compress_ratios_sorted=ratios, + max_draft_tokens=5, + past_kv_lens_cuda={ratio: torch.empty(8, dtype=torch.int32) for ratio in ratios}, + cu_new_comp_kv_cuda={ratio: torch.empty(9, dtype=torch.int32) for ratio in ratios}, + new_comp_kv_lens_cuda={ratio: torch.empty(8, dtype=torch.int32) for ratio in ratios}, + compressed_position_ids_cuda={ + ratio: torch.empty(48, dtype=torch.int32) for ratio in ratios + }, + compressed_mask_cuda={ratio: torch.empty(48, dtype=torch.bool) for ratio in ratios}, + batch_size_cuda=torch.tensor([17], dtype=torch.int32), + _compute_gen_compressed_position_ids=position_helper, + _compute_compressed_mask=mask_helper, + ) + + +def test_ragged_compressor_warmup_primes_and_cleans_intermediate_shape(monkeypatch): + position_calls = [] + mask_calls = [] + + def record_positions(*args): + position_calls.append((torch.is_inference_mode_enabled(), args[3:])) + + def record_mask(new_comp, cu_new_comp, compressed_mask, batch_size, compress_ratios): + mask_calls.append((torch.is_inference_mode_enabled(), batch_size.tolist(), compress_ratios)) + + metadata = _warmup_metadata(record_positions, record_mask) + engine = SimpleNamespace( + _ragged_verification_enabled=True, + attn_metadata=metadata, + _config=SimpleNamespace(max_batch_size=8, cuda_graph_batch_sizes=[1, 4, 8]), + ) + sync_calls = [] + monkeypatch.setattr(torch.cuda, "synchronize", lambda: sync_calls.append(True)) + + DecoderMetadata._warmup_dspark_ragged_compressor_metadata(engine) + + assert position_calls == [(True, (0, 7, 6, [1, 4], {1: 0, 4: 0}))] + assert mask_calls == [(True, [7], [1, 4])] + assert metadata.batch_size_cuda.tolist() == [17] + for ratio, token_count in ((1, 42), (4, 14)): + assert not metadata.past_kv_lens_cuda[ratio][:7].any() + assert not metadata.new_comp_kv_lens_cuda[ratio][:7].any() + assert not metadata.cu_new_comp_kv_cuda[ratio][:8].any() + assert not metadata.compressed_position_ids_cuda[ratio][:token_count].any() + assert not metadata.compressed_mask_cuda[ratio][:token_count].any() + assert sync_calls == [True, True] + + +def test_ragged_compressor_warmup_cleans_after_helper_failure(monkeypatch): + def fail_positions(*args): + args[2][1][:18].fill_(7) + raise RuntimeError("compile boom") + + metadata = _warmup_metadata(fail_positions, lambda *_args: None) + engine = SimpleNamespace( + _ragged_verification_enabled=True, + attn_metadata=metadata, + _config=SimpleNamespace(max_batch_size=4, cuda_graph_batch_sizes=[1, 4]), + ) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + + with pytest.raises(RuntimeError, match="compile boom"): + DecoderMetadata._warmup_dspark_ragged_compressor_metadata(engine) + + assert not metadata.past_kv_lens_cuda[1][:3].any() + assert not metadata.new_comp_kv_lens_cuda[1][:3].any() + assert not metadata.cu_new_comp_kv_cuda[1][:4].any() + assert not metadata.compressed_position_ids_cuda[1][:18].any() + assert not metadata.compressed_mask_cuda[1][:18].any() + assert metadata.batch_size_cuda.tolist() == [17] + + +@pytest.mark.parametrize("fail_after_mask", [False, True]) +def test_ragged_compressor_warmup_calls_current_metadata_helpers(monkeypatch, fail_after_mask): + """Run the real five-argument mask and position helper bodies on CPU tensors.""" + actual_mask = DeepseekV4TrtllmAttentionMetadata._compute_compressed_mask + + def mask(new_comp, cu_new_comp, compressed_mask, batch_size, compress_ratios): + actual_mask(new_comp, cu_new_comp, compressed_mask, batch_size, compress_ratios) + assert batch_size is metadata.batch_size_cuda + assert batch_size.tolist() == [7] + assert compressed_mask[1][:42].all() + assert compressed_mask[4][:14].all() + assert metadata.compressed_position_ids_cuda[1][:42].tolist() == list(range(6)) * 7 + assert metadata.compressed_position_ids_cuda[4][:14].tolist() == [0, 4] * 7 + if fail_after_mask: + raise RuntimeError("mask cleanup probe") + + metadata = _warmup_metadata( + DeepseekV4TrtllmAttentionMetadata._compute_gen_compressed_position_ids, mask + ) + original_storage = metadata.batch_size_cuda.data_ptr() + engine = SimpleNamespace( + _ragged_verification_enabled=True, + attn_metadata=metadata, + _config=SimpleNamespace(max_batch_size=8, cuda_graph_batch_sizes=[1, 4, 8]), + ) + monkeypatch.setattr(torch_utils, "is_piecewise_running", lambda: True) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + if fail_after_mask: + with pytest.raises(RuntimeError, match="mask cleanup probe"): + DecoderMetadata._warmup_dspark_ragged_compressor_metadata(engine) + else: + DecoderMetadata._warmup_dspark_ragged_compressor_metadata(engine) + assert metadata.batch_size_cuda.data_ptr() == original_storage + assert metadata.batch_size_cuda.tolist() == [17] + for ratio, token_count in ((1, 42), (4, 14)): + assert not metadata.past_kv_lens_cuda[ratio][:7].any() + assert not metadata.new_comp_kv_lens_cuda[ratio][:7].any() + assert not metadata.cu_new_comp_kv_cuda[ratio][:8].any() + assert not metadata.compressed_position_ids_cuda[ratio][:token_count].any() + assert not metadata.compressed_mask_cuda[ratio][:token_count].any() + + +def test_ragged_compressor_warmup_requires_trimmed_tokens(): + engine = SimpleNamespace( + _ragged_verification_enabled=False, + ) + + DecoderMetadata._warmup_dspark_ragged_compressor_metadata(engine) + + +def test_ragged_compressor_warmup_rejects_missing_metadata_contract(): + engine = SimpleNamespace( + _ragged_verification_enabled=True, + attn_metadata=SimpleNamespace(), + ) + + with pytest.raises(RuntimeError, match="missing .*max_draft_tokens"): + DecoderMetadata._warmup_dspark_ragged_compressor_metadata(engine) + + +@pytest.mark.parametrize("runtime_draft_len", [0, 2, 5]) +def test_publish_gen_token_layout_uses_call_draft_length(runtime_draft_len: int) -> None: + engine = _engine(_runner(), []) + engine.get_runtime_tokens_per_gen_step = lambda draft_len: draft_len + 1 + requests = [ + SimpleNamespace(py_verify_len=min(1, runtime_draft_len)), + SimpleNamespace(py_verify_len=min(3, runtime_draft_len)), + ] + metadata = SimpleNamespace( + runtime_tokens_per_gen_step=0, + ragged_verify_lens=None, + device_windows_mode=False, + ) + assert not hasattr(engine, "runtime_draft_len") + + engine._publish_gen_token_layout(metadata, requests, runtime_draft_len=runtime_draft_len) + + assert metadata.runtime_tokens_per_gen_step == runtime_draft_len + 1 + assert metadata.ragged_verify_lens == [request.py_verify_len + 1 for request in requests] + assert not metadata.device_windows_mode + assert not hasattr(engine, "runtime_draft_len") + + +def test_attach_ragged_layout_refreshes_persistent_buffers_in_place() -> None: + runner = _engine(_runner(), []) + runner._pinned_host_cache = {} + runner._pinned_host_active = {} + runner._buffers = SimpleNamespace() + runner._buffers.ragged_verify_lens_cuda = torch.empty(9, dtype=torch.int32) + runner._buffers.ragged_qo_indptr_cuda = torch.empty(10, dtype=torch.int32) + attention = SimpleNamespace(ragged_verify_lens=None) + speculation = SimpleNamespace() + requests = [SimpleNamespace(py_verify_len=1), SimpleNamespace(py_verify_len=3)] + + runner._attach_ragged_verify_layout(speculation, attention, requests) + first_lens_ptr = speculation.verify_lens.data_ptr() + first_prefix_ptr = speculation.qo_indptr.data_ptr() + + assert speculation.verify_lens.tolist() == [2, 4] + assert speculation.qo_indptr.tolist() == [0, 2, 6] + assert speculation.total_verify_tokens == 6 + assert attention.ragged_verify_lens == [2, 4] + + requests[0].py_verify_len = 0 + requests[1].py_verify_len = 2 + runner._attach_ragged_verify_layout(speculation, attention, requests) + + assert speculation.verify_lens.tolist() == [1, 3] + assert speculation.qo_indptr.tolist() == [0, 1, 4] + assert speculation.total_verify_tokens == 4 + assert attention.ragged_verify_lens == [1, 3] + assert speculation.verify_lens.data_ptr() == first_lens_ptr + assert speculation.qo_indptr.data_ptr() == first_prefix_ptr + + +@pytest.mark.parametrize("trim_enabled", [False, True]) +def test_missing_or_disabled_windows_clear_stale_ragged_views(trim_enabled: bool) -> None: + runner = _engine(_runner(), []) + runner._dspark_trims_submitted_tokens = trim_enabled + attention = SimpleNamespace(ragged_verify_lens=[2, 4]) + speculation = SimpleNamespace( + verify_lens=torch.tensor([2, 4]), + qo_indptr=torch.tensor([0, 2, 6]), + total_verify_tokens=6, + ) + + runner._attach_ragged_verify_layout( + speculation, + attention, + [SimpleNamespace(py_verify_len=1), SimpleNamespace(py_verify_len=None)], + ) + + assert speculation.verify_lens is None + assert speculation.qo_indptr is None + assert speculation.total_verify_tokens is None + assert attention.ragged_verify_lens is None + + +def test_model_engine_delegates_ragged_state_to_decoder_owner() -> None: + inputs = _engine(_runner(), []) + runner = object.__new__(DecoderRunner) + runner._inputs = inputs + inputs._dspark_sps_cost_table = object() + inputs._dspark_exact_candidate_cells = ((4, 12),) + inputs._dspark_exact_identity_words = tuple(range(8)) + inputs.fit_ragged_verify_lens = lambda *args, **kwargs: (args, kwargs) + worker = object() + runner.model = SimpleNamespace(spec_worker=worker) + engine = object.__new__(PyTorchModelEngine) + engine._cleanup_done = True + engine._runner = runner + + assert engine._dspark_confidence_enabled + assert engine._dspark_trims_submitted_tokens + assert engine._dspark_sps_cost_table is runner._dspark_sps_cost_table + assert engine._dspark_exact_candidate_cells == ((4, 12),) + assert engine._dspark_exact_identity_words == tuple(range(8)) + assert engine._dspark_device_budget is None + assert engine._get_spec_worker() is worker + engine._dspark_device_budget = 7 + assert runner._dspark_device_budget == 7 + assert engine.fit_ragged_verify_lens([1], exact_shape=(4, 12, 3)) == ( + ([1],), + {"exact_shape": (4, 12, 3)}, + ) + + +def test_model_engine_non_decoder_keeps_confidence_inactive() -> None: + engine = object.__new__(PyTorchModelEngine) + engine._cleanup_done = True + engine._runner = SimpleNamespace() + + assert not engine._dspark_confidence_enabled + assert not engine._dspark_trims_submitted_tokens + assert engine._dspark_sps_cost_table is None + assert engine._dspark_exact_candidate_cells == () + assert engine._dspark_exact_identity_words == (0,) * 8 + assert engine._dspark_device_budget is None + assert engine._get_spec_worker() is None + + +def _metadata_inputs(metadata_type): + model = SimpleNamespace( + pretrained_config=SimpleNamespace(num_attention_heads=4, num_key_value_heads=1), + sparse_attention_config=None, + enable_flash_mla=False, + ) + kwargs = dict( + max_batch_size=2, + max_num_tokens=8, + max_beam_width=1, + attention_backend=SimpleNamespace(Metadata=metadata_type), + attention_runtime_features=SimpleNamespace(cache_reuse=False, chunked_prefill=False), + mapping=object(), + cache_indirection=None, + enable_context_mla_with_cached_kv=False, + ) + return model, kwargs + + +@pytest.mark.parametrize("explicit_false", [False, True]) +def test_metadata_helper_keeps_native_constructor_kwargs(explicit_false): + seen = [] + + class NativeMetadata: + def __init__(self, **kwargs): + seen.append(kwargs) + + model, kwargs = _metadata_inputs(NativeMetadata) + if explicit_false: + kwargs["enable_ragged_verification"] = False + metadata_module.build_attention_metadata(model, **kwargs) + + assert "enable_ragged_verification" not in seen[0] + assert seen[0]["num_heads_per_kv"] == 4 + + +def test_metadata_helper_enables_only_supported_ragged_metadata(): + class RaggedMetadata(DSAtrtllmAttentionMetadata): + def __init__(self, **kwargs): + self.constructor_flag = kwargs.pop("enable_ragged_verification", False) + self.enable_ragged_verification = self.constructor_flag + self.ragged_verify_lens = [1, 2] + + model, kwargs = _metadata_inputs(RaggedMetadata) + result = metadata_module.build_attention_metadata( + model, enable_ragged_verification=True, **kwargs + ) + assert result.constructor_flag is True + result._validate_ragged_verification() + + kwargs["attention_backend"] = SimpleNamespace(Metadata=object) + with pytest.raises(ValueError, match="DSA target metadata"): + metadata_module.build_attention_metadata(model, enable_ragged_verification=True, **kwargs) + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_decoder_enables_ragged_metadata_during_construction(monkeypatch, enabled): + seen = [] + + class MetadataProbe(DSAtrtllmAttentionMetadata): + def __init__(self, **kwargs): + self.enable_ragged_verification = kwargs.pop("enable_ragged_verification", False) + self.kv_cache_manager = kwargs["kv_cache_manager"] + self.ragged_verify_lens = None + seen.append(self.enable_ragged_verification) + + model, kwargs = _metadata_inputs(MetadataProbe) + decoder = object.__new__(DecoderMetadata) + decoder.attn_metadata = None + decoder.model = SimpleNamespace(model_config=model) + decoder._config = SimpleNamespace( + **kwargs, spec_config=SimpleNamespace(enable_confidence_scheduling=enabled) + ) + decoder.mapping = kwargs["mapping"] + decoder._buffers = SimpleNamespace(cache_indirection=None) + monkeypatch.setattr(decoder_metadata_module, "resolve_mamba_metadata_cls", lambda _model: None) + kv = object() + + result = decoder.setup_attention(kv) + if enabled: + result.ragged_verify_lens = [1, 2] + result._validate_ragged_verification() + + assert seen == [enabled] + assert decoder.setup_attention(kv) is result + assert seen == [enabled] + + +@pytest.mark.parametrize("graph_type", [CUDAGraphRunner, DSparkCUDAGraphRunner]) +def test_current_main_warmup_keeps_serving_timeout_only_for_capture(monkeypatch, graph_type): + graph = graph_type.__new__(graph_type) + graph.clear = lambda: None + graph._capture_allowed = False + graph.padding_dummy_requests = {5: object()} + if graph_type is DSparkCUDAGraphRunner: + graph.secondary_padding_dummy_requests = {5: object()} + graph.ragged_zero_real_high_rows = 2 + graph.adp_shape_agreement = object() + serving_phase = False + calls = [] + + @contextmanager + def serving_timeouts(): + nonlocal serving_phase + serving_phase = True + try: + yield + finally: + serving_phase = False + + def warmup(_resources): + calls.append((graph.is_warmup_only, serving_phase, bool(graph.padding_dummy_requests))) + assert graph._capture_allowed + if not graph.is_warmup_only and graph_type is DSparkCUDAGraphRunner: + assert graph.secondary_padding_dummy_requests == {} + assert graph.ragged_zero_real_high_rows == 0 + assert graph.adp_shape_agreement is None + + monkeypatch.setattr(graphs_module, "moe_comm_serving_timeouts", serving_timeouts) + decoder = object.__new__(DecoderCudaGraphs) + decoder.cuda_graph_runner = graph + decoder.run_cuda_graph_warmup = warmup + decoder.warmup_and_capture(object()) + + assert calls == [(True, False, True), (False, True, False)] + assert not graph._capture_allowed + assert not graph.is_warmup_only + + +@pytest.mark.parametrize("worker", [None, SimpleNamespace(verify_planner=None)]) +def test_device_window_prologue_declines_an_unavailable_planner( + worker: SimpleNamespace | None, +) -> None: + decoder = object.__new__(DecoderInputPreparer) + decoder.model = SimpleNamespace(spec_worker=worker) + decoder.cuda_graph_runner = SimpleNamespace(agreed_ragged_bucket=10) + decoder._dspark_device_budget = 10 + decoder._dspark_prev_covers_batch = True + + assert not decoder._apply_device_window_prologue({}, None) + assert decoder._dspark_device_budget is None + + +def test_confidence_without_an_sps_table_is_eager_only() -> None: + decoder = object.__new__(DecoderInputPreparer) + decoder.cuda_graph_runner = SimpleNamespace(enabled=False) + decoder._dspark_confidence_enabled = True + decoder._dspark_trims_submitted_tokens = True + decoder._dspark_sps_cost_table = None + decoder._config = SimpleNamespace( + spec_config=SimpleNamespace(draft_len_schedule=None, max_draft_len=3) + ) + + graphs = object.__new__(DecoderCudaGraphs) + graphs._inputs = decoder + graphs._config = decoder._config + graphs.cuda_graph_runner = decoder.cuda_graph_runner + graphs._capture_generation_cuda_graphs(None) + with pytest.raises(RuntimeError, match="authenticated exact SPS table"): + graphs._get_graphs_to_capture([1, 4]) + + +@pytest.mark.parametrize( + ("trims", "prefix_window", "previous_windows"), + [(True, 1, [3, 3]), (True, 1, [1, 2]), (False, None, [None, None])], +) +def test_overlap_offsets_follow_the_packed_prefix( + trims: bool, prefix_window: int | None, previous_windows: list[int | None] +) -> None: + """Exercise the full input packer with CPU buffers and metadata-only leaves.""" + decoder = object.__new__(DecoderInputPreparer) + mode = SimpleNamespace( + support_overlap_scheduler=lambda: True, + extend_ctx=lambda _backend: False, + use_one_engine=lambda: True, + ) + spec_config = SimpleNamespace(spec_dec_mode=mode, is_linear_tree=True, max_draft_len=3) + decoder._config = SimpleNamespace( + use_mrope=False, + max_num_tokens=32, + max_beam_width=1, + is_spec_decode=True, + disable_overlap_scheduler=False, + spec_config=spec_config, + attention_backend="TRTLLM", + enable_attention_dp=False, + prefill_cuda_graph_backend=decoder_module.PrefillCudaGraphBackend.DISABLED, + prefill_cuda_graph_num_tokens=None, + ) + decoder.mapping = SimpleNamespace(pp_size=1, cp_config={}, has_cp_helix=lambda: False) + decoder.dist = None + decoder.guided_decoder = None + decoder.model = SimpleNamespace( + model_config=SimpleNamespace(pretrained_config=SimpleNamespace()) + ) + decoder._dspark_trims_submitted_tokens = trims + decoder._dspark_device_windows = False + decoder._prepare_inputs_fast_path = lambda *_args, **_kwargs: None + decoder._can_use_steady_gen_fast_prepare = lambda *_args: False + decoder._new_extra_inputs = lambda: SimpleNamespace( + add_generation_request=lambda _request: None, build=lambda *_args: {} + ) + decoder._lora = SimpleNamespace(build=lambda *_args, **_kwargs: None) + decoder._should_use_full_generation_page_table = lambda *_args: False + decoder.get_runtime_tokens_per_gen_step = lambda draft_len: draft_len + 1 + decoder.iter_states = {} + decoder.log_cached_kv_tokens_per_req = False + decoder.use_beam_search = False + decoder._mm_item_scheduler = None + decoder._pinned_host_cache = {} + decoder._pinned_host_active = {} + decoder._ragged_gather_indices = lambda slots, lengths: tuple( + torch.tensor(indices, dtype=torch.long) + for indices in ragged_gather_index_lists(slots, lengths) + ) + decoder._buffers = SimpleNamespace() + for name in ( + "input_ids_cuda", + "draft_tokens_cuda", + "num_accepted_draft_tokens_cuda", + "previous_pos_id_offsets_cuda", + "previous_kv_lens_offsets_cuda", + "previous_pos_indices_cuda", + "previous_batch_indices_cuda", + "position_ids_cuda", + "gather_ids_cuda", + ): + setattr(decoder._buffers, name, torch.full((32,), -99, dtype=torch.int32)) + windows = [prefix_window, *previous_windows] + requests = [ + SimpleNamespace( + py_request_id=index + 100, + py_seq_slot=index, + py_batch_idx=None if index == 0 else index, + py_draft_tokens=[11, 12, 13], + py_verify_len=window, + is_dummy=False, + is_attention_dp_dummy=False, + is_cuda_graph_dummy=False, + py_is_first_draft=False, + py_num_accepted_draft_tokens=1, + max_beam_num_tokens=8, + py_prompt_len=4, + py_num_compressed_tokens=0, + get_last_tokens=lambda _beam: 10, + ) + for index, window in enumerate(windows) + ] + scheduled = SimpleNamespace( + context_requests=[], + generation_requests=requests, + num_context_requests=0, + num_generation_requests=3, + is_attention_dp_phase_idle=False, + attention_dp_phase=None, + ) + state = object.__new__(decoder_module.SampleStateTensorsSpec) + state.new_tokens = torch.arange(12, dtype=torch.int32).reshape(4, 3) + state.new_tokens_lens = torch.tensor([2, 3, 4], dtype=torch.int32) + state.next_draft_tokens = torch.arange(9, dtype=torch.int32).reshape(3, 3) + metadata = SimpleNamespace(is_cuda_graph=False, prepare=lambda: None) + + prepared = decoder.prepare_inputs( + scheduled, + None, + metadata, + new_tensors_device=state, + enable_spec_decode=True, + runtime_draft_len=3, + is_dummy=False, + ) + + inputs, gather_ids, draft_len = ( + prepared.kwargs, + prepared.gather_ids, + prepared.runtime_draft_len, + ) + lengths = [1 + window if trims else 4 for window in windows] + prefix_len, *previous_lengths = lengths + expected = [0] * prefix_len + [3] * previous_lengths[0] + [4] * previous_lengths[1] + assert decoder._buffers.previous_pos_id_offsets_cuda[: sum(lengths)].tolist() == expected + assert decoder._buffers.previous_pos_id_offsets_cuda[sum(lengths) :].count_nonzero().item() == 0 + assert decoder._buffers.previous_kv_lens_offsets_cuda[:3].tolist() == [ + 0, + 3 - previous_lengths[0], + 4 - previous_lengths[1], + ] + assert inputs["input_ids"].numel() == sum(lengths) + assert metadata.seq_lens.tolist() == lengths + assert metadata.kv_cache_params.num_extra_kv_tokens == 2 + assert gather_ids.tolist() == list(range(sum(lengths))) + assert draft_len == 3 diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_select.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_select.py new file mode 100644 index 000000000000..0d1d9d737417 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_select.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU tests for fresh-confidence DSpark device-window selection.""" + +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig +from tensorrt_llm._torch.speculative.dspark_select import ( + gather_packed_draft_tokens, + select_windows_device, +) + + +def test_fresh_confidence_spends_budget_on_the_best_real_row(): + result = select_windows_device( + confidence_logits=torch.tensor( + [[10.0, 10.0, 10.0, 10.0, 10.0], [-10.0, -10.0, -10.0, -10.0, -10.0]] + ), + slot_idx=torch.tensor([0, 1, 0, 0]), + num_real=torch.tensor(2), + budget=torch.tensor(2), + graph_num_tokens=8, + cfg=DSparkScheduleConfig(block_size=5, min_verify_len=1), + pad_len=1, + ) + + assert result.verify_lens.tolist() == [4, 2, 1, 1] + assert result.qo_indptr.tolist() == [0, 4, 6, 7, 8] + assert result.req_idx.tolist() == [0, 0, 0, 0, 1, 1, 2, 3] + assert result.kv_correction.tolist() == [-3, -2, -1, 0, -1, 0, 0, 0] + + +def test_stale_row_fails_open_to_neutral_confidence(): + result = select_windows_device( + confidence_logits=torch.full((2, 5), -10.0), + slot_idx=torch.tensor([0, 1]), + num_real=torch.tensor(2), + budget=torch.tensor(1), + graph_num_tokens=5, + cfg=DSparkScheduleConfig(block_size=5, min_verify_len=1), + stamp=torch.tensor([7, 6]), + expected_stamp=torch.tensor(7), + ) + + assert result.verify_lens.tolist() == [2, 3] + + +def test_gather_packed_draft_tokens_excludes_each_bonus(): + packed = gather_packed_draft_tokens( + next_draft_tokens=torch.tensor([[10, 11, 12], [20, 21, 22]]), + batch_slots=torch.tensor([1, 0]), + verify_lens=torch.tensor([3, 2], dtype=torch.int32), + qo_indptr=torch.tensor([0, 3, 5], dtype=torch.int32), + num_real=2, + total_draft_tokens=3, + ) + + assert packed.tolist() == [20, 21, 10] + + +def test_gather_packed_draft_tokens_rejects_negative_size(): + with pytest.raises(ValueError, match="non-negative"): + gather_packed_draft_tokens( + next_draft_tokens=torch.empty((1, 1)), + batch_slots=torch.tensor([0]), + verify_lens=torch.tensor([1], dtype=torch.int32), + qo_indptr=torch.tensor([0, 1], dtype=torch.int32), + num_real=1, + total_draft_tokens=-1, + ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index 40786627b200..6a978a5aeb08 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -136,8 +136,8 @@ def test_worker_lazy_init_window_buffers(): assert worker._valid_len.shape == (9,) assert worker._position_initialized.shape == (9,) assert worker._scratch_slot == 8 - # Dummy-id floor separates real request ids from CUDA-graph padding ids. - assert worker._graph_dummy_id_floor == (1 << 64) - 1 - worker.max_draft_len + # Both CUDA-graph padding families lie at or above the dummy-id floor. + assert worker._graph_dummy_id_floor == (1 << 64) - 2 * (worker.max_draft_len + 1) # The scratch row is never handed out through the free pool. assert list(worker._free_slots) == list(range(8)) assert worker._batch_to_slot is not None diff --git a/tests/unittest/_torch/speculative/test_device_window_sampling_layout.py b/tests/unittest/_torch/speculative/test_device_window_sampling_layout.py new file mode 100644 index 000000000000..4584b10b544a --- /dev/null +++ b/tests/unittest/_torch/speculative/test_device_window_sampling_layout.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import types + +import pytest +import torch + +from tensorrt_llm._torch.speculative.interface import _RNG_SLOT_SPAN, SpecMetadata + +pytestmark = pytest.mark.cpu_only + + +def test_device_window_repacking_follows_fresh_token_owners(): + owners = torch.tensor([0, 0, 1, 2, 2, 2], dtype=torch.long) + meta = types.SimpleNamespace( + request_temperatures=torch.tensor([0.2, 0.7, 1.1]), + temperatures=torch.empty(6), + request_top_ks=torch.tensor([2, 7, 11], dtype=torch.int32), + top_ks=torch.empty(6, dtype=torch.int32), + request_top_ps=torch.tensor([0.3, 0.8, 0.95]), + top_ps=torch.empty(6), + request_min_ps=torch.tensor([0.1, 0.2, 0.4]), + min_ps=torch.empty(6), + request_seeds=torch.tensor([101, 202, 303], dtype=torch.int64), + seeds=torch.empty(6, dtype=torch.int64), + request_offsets=torch.tensor([10, 20, 30], dtype=torch.int64), + offsets=torch.empty(6, dtype=torch.int64), + qo_indptr=torch.tensor([0, 2, 3, 6], dtype=torch.int32), + _sampling_params_signature=[("request",), ("expanded",)], + ) + + SpecMetadata.remap_expanded_sampling_params(meta, owners, owners.numel()) + + assert torch.equal(meta.temperatures, meta.request_temperatures.index_select(0, owners)) + assert torch.equal(meta.top_ks, meta.request_top_ks.index_select(0, owners)) + assert torch.equal(meta.top_ps, meta.request_top_ps.index_select(0, owners)) + assert torch.equal(meta.min_ps, meta.request_min_ps.index_select(0, owners)) + assert torch.equal(meta.seeds, meta.request_seeds.index_select(0, owners)) + row_slots = torch.tensor([0, 1, 0, 0, 1, 2], dtype=torch.int64) + expected_offsets = meta.request_offsets.index_select(0, owners) + row_slots * _RNG_SLOT_SPAN + assert torch.equal(meta.offsets, expected_offsets) + assert meta._sampling_params_signature[0] == ("request",) + assert meta._sampling_params_signature[1] is None + + +def test_device_window_repacking_ignores_unallocated_optional_buffers(): + meta = types.SimpleNamespace( + request_temperatures=None, + temperatures=None, + request_top_ks=None, + top_ks=None, + request_top_ps=None, + top_ps=None, + request_min_ps=None, + min_ps=None, + request_seeds=None, + seeds=None, + request_offsets=None, + offsets=None, + _sampling_params_signature=[None, ("expanded",)], + ) + SpecMetadata.remap_expanded_sampling_params(meta, torch.tensor([0]), 1) + assert meta._sampling_params_signature[1] is None + + +@pytest.mark.parametrize("missing_row_starts", [False, True]) +def test_device_window_repacking_rejects_invalid_captured_buffers(missing_row_starts): + meta = types.SimpleNamespace( + request_temperatures=torch.tensor([0.2]), + temperatures=torch.empty(2 if missing_row_starts else 1), + request_top_ks=None, + top_ks=None, + request_top_ps=None, + top_ps=None, + request_min_ps=None, + min_ps=None, + request_seeds=None, + seeds=None, + request_offsets=torch.tensor([10], dtype=torch.int64) if missing_row_starts else None, + offsets=torch.empty(2, dtype=torch.int64) if missing_row_starts else None, + qo_indptr=None, + _sampling_params_signature=[None, ("expanded",)], + ) + expected_error = "ragged row starts" if missing_row_starts else "exceeds the captured" + with pytest.raises(RuntimeError, match=expected_error): + SpecMetadata.remap_expanded_sampling_params(meta, torch.tensor([0, 0]), 2) + + +def test_device_window_repacking_forces_next_host_expansion(): + values = ((0.2, 2, 0.3, 0.1), (0.7, 7, 0.8, 0.2)) + per_request = [(0.2, 2, 0.3, 0.1, 2), (0.7, 7, 0.8, 0.2, 4)] + meta = types.SimpleNamespace( + request_temperatures=torch.tensor([0.2, 0.7]), + temperatures=torch.empty(6), + request_top_ks=torch.tensor([2, 7], dtype=torch.int32), + top_ks=torch.empty(6, dtype=torch.int32), + request_top_ps=torch.tensor([0.3, 0.8]), + top_ps=torch.empty(6), + request_min_ps=torch.tensor([0.1, 0.2]), + min_ps=torch.empty(6), + request_seeds=None, + seeds=None, + request_offsets=None, + offsets=None, + _sampling_params_signature=[values, (values, (2, 4))], + ) + SpecMetadata.remap_expanded_sampling_params(meta, torch.tensor([0, 1, 1, 1, 0, 0]), 6) + + need_request, need_expanded = SpecMetadata._sampling_params_buffers_need_update( + meta, per_request + ) + assert need_request is False + assert need_expanded is True From eb8abd228bf5c7c162864cbdfdb44b846533db2e Mon Sep 17 00:00:00 2001 From: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> Date: Mon, 19 Oct 2026 18:26:15 -0700 Subject: [PATCH 2/2] [TRTLLM-14620][fix] Preserve decoder construction contracts Keep input-preparer ownership in the native and encoder-decoder graph factories, and retain optional-config fixture construction without changing configured confidence behavior. Add focused factory and constructor regression coverage. Signed-off-by: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> --- .../engine/runners/decoder/inputs.py | 44 ++++--- .../engine/runners/decoder/runner.py | 3 +- .../executor/test_pytorch_model_engine.py | 87 +++++++++++-- .../test_dspark_cuda_graph_runtime.py | 6 +- .../test_dspark_model_engine_runtime.py | 120 +++++++++++++++++- 5 files changed, 219 insertions(+), 41 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py index 292344be85a6..74142811f4f2 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py @@ -357,7 +357,7 @@ def __init__( self._prepare_inputs_event: torch.cuda.Event | None = None self._gather_ids_buffer = ( DeviceCopyBuffer(buffers.gather_ids_cuda, pin_memory=prefer_pinned()) - if config.is_spec_decode + if config is not None and config.is_spec_decode else None ) # Reuse device indices while the scheduled batch stays unchanged. @@ -365,7 +365,7 @@ def __init__( self._staged_previous_pos_indices: list[int] | None = None self._step_input_gather: spec_step_copies.StepInputGather | None = ( spec_step_copies.StepInputGather() - if config.is_spec_decode and spec_step_copies.is_supported() + if config is not None and config.is_spec_decode and spec_step_copies.is_supported() else None ) # Steady-state generation-only prepare cache (non-speculative overlap @@ -381,7 +381,7 @@ def __init__( os.getenv("TLLM_LOG_CACHED_KV_TOKENS_PER_REQ", "0") == "1" ) - spec_config = self._config.spec_config + spec_config = self._config.spec_config if self._config is not None else None self._dspark_confidence_enabled = bool( spec_config is not None and getattr(spec_config, "enable_confidence_scheduling", False) ) @@ -1950,23 +1950,27 @@ def previous_seq_slots_device() -> torch.Tensor: if host.previous_batch_token_start is not None else previous_begin * runtime_tokens_per_gen_step ) - if is_ragged_gen or step_gather is None or not step_gather.gather( - new_tokens_device, - next_draft_tokens_device, - new_tokens_lens_device, - self._buffers.previous_batch_indices_cuda, - self._buffers.previous_pos_indices_cuda, - previous_batch_len, - runtime_tokens_per_gen_step, - runtime_draft_token_buffer_width, - self._buffers.input_ids_cuda, - num_tokens, - self._buffers.draft_tokens_cuda, - num_draft_tokens, - self._buffers.previous_pos_id_offsets_cuda, - pos_offsets_start, - self._buffers.previous_kv_lens_offsets_cuda, - previous_begin, + if ( + is_ragged_gen + or step_gather is None + or not step_gather.gather( + new_tokens_device, + next_draft_tokens_device, + new_tokens_lens_device, + self._buffers.previous_batch_indices_cuda, + self._buffers.previous_pos_indices_cuda, + previous_batch_len, + runtime_tokens_per_gen_step, + runtime_draft_token_buffer_width, + self._buffers.input_ids_cuda, + num_tokens, + self._buffers.draft_tokens_cuda, + num_draft_tokens, + self._buffers.previous_pos_id_offsets_cuda, + pos_offsets_start, + self._buffers.previous_kv_lens_offsets_cuda, + previous_begin, + ) ): # previous input ids if is_ragged_gen: diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 391fc4d50386..49ae6f0c6cb7 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -232,7 +232,7 @@ def _make_input_preparer( def _make_cuda_graphs(self) -> DecoderCudaGraphs: """Build the CUDA graph component; call after the input preparer.""" - return DecoderCudaGraphs(**self._cuda_graphs_kwargs()) + return DecoderCudaGraphs(**self._cuda_graphs_kwargs(), inputs=self._inputs) def _dummy_requests_kwargs(self) -> dict[str, Any]: return dict( @@ -281,7 +281,6 @@ def _cuda_graphs_kwargs(self) -> dict[str, Any]: torch_compile_bypass=self._bypass_torch_compile, context_promotion_supported=self._context_graph_promotion_supported, warmup_inputs=self._buffers.input_ids_cuda, - inputs=self._inputs, ) @property diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 2fb36454174a..c18470e44f87 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -39,6 +39,8 @@ PrefillCompiledModel from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.warmup import \ DecoderWarmup +from tensorrt_llm._torch.pyexecutor.engine.runners.encoder_decoder import \ + EncoderDecoderRunner from tensorrt_llm._torch.pyexecutor.engine.runners.interface import \ PreparedInputs from tensorrt_llm._torch.pyexecutor.engine.runners.pooling import PoolingRunner @@ -78,17 +80,80 @@ _TWO_MODEL_SPEC_MODE = SimpleNamespace(use_one_engine=lambda: False) -def _input_preparer() -> DecoderInputPreparer: - return DecoderInputPreparer(buffers=None, - config=SimpleNamespace(is_spec_decode=False), - mapping=None, - dist=None, - model=None, - get_runtime_tokens_per_gen_step=None, - iter_states={}, - lora=None, - use_beam_search=False, - steady_gen_cache_supported=True) +def _input_preparer(*, without_config: bool = False) -> DecoderInputPreparer: + return DecoderInputPreparer( + buffers=None, + config=(None if without_config else SimpleNamespace( + is_spec_decode=False, spec_config=None)), + mapping=None, + dist=None, + model=None, + get_runtime_tokens_per_gen_step=None, + iter_states={}, + lora=None, + use_beam_search=False, + steady_gen_cache_supported=True) + + +@pytest.mark.parametrize("runner_type,graphs_path", [ + (DecoderRunner, + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner.DecoderCudaGraphs" + ), + (EncoderDecoderRunner, + "tensorrt_llm._torch.pyexecutor.engine.runners.encoder_decoder.EncoderDecoderCudaGraphs" + ) +]) +def test_cuda_graph_factory_preserves_input_preparer(runner_type, graphs_path): + runner = Mock() + runner._inputs = object() + runner._cuda_graphs_kwargs.side_effect = lambda: DecoderRunner._cuda_graphs_kwargs( + runner) + + with patch(graphs_path) as graphs_type: + graphs = runner_type._make_cuda_graphs(runner) + + graphs_type.assert_called_once() + assert graphs_type.call_args.kwargs["inputs"] is runner._inputs + assert graphs_type.call_args.kwargs[ + "warmup_inputs"] is runner._buffers.input_ids_cuda + assert graphs is graphs_type.return_value + + +@pytest.mark.parametrize("enabled,fused", [(None, False), (False, True), + (True, False), (True, True)]) +def test_input_preparer_configured_confidence_flags(enabled, fused): + spec_config = (None if enabled is None else SimpleNamespace( + enable_confidence_scheduling=enabled, + enable_fused_confidence_scheduler=fused)) + config = SimpleNamespace(is_spec_decode=False, spec_config=spec_config) + inputs = DecoderInputPreparer(buffers=None, + config=config, + mapping=None, + dist=None, + model=None, + get_runtime_tokens_per_gen_step=None, + iter_states={}, + lora=None, + use_beam_search=False, + steady_gen_cache_supported=True) + + assert inputs._config is config + assert inputs._dspark_confidence_enabled is bool(enabled) + assert inputs._dspark_trims_submitted_tokens is bool(enabled) + assert inputs._dspark_device_windows is bool(enabled and fused) + assert inputs.previous_verify_lens == ([] if enabled else None) + assert inputs.has_previous_device_draft is False + + +def test_input_preparer_without_config_keeps_confidence_disabled(): + inputs = _input_preparer(without_config=True) + + assert inputs._config is None + assert inputs._dspark_confidence_enabled is False + assert inputs._dspark_trims_submitted_tokens is False + assert inputs._dspark_device_windows is False + assert inputs.previous_verify_lens is None + assert inputs.has_previous_device_draft is False class TestCachedKvTokenLogging(unittest.TestCase): diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph_runtime.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph_runtime.py index fb5e02160175..dd59b0292daa 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph_runtime.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph_runtime.py @@ -42,7 +42,7 @@ def num_generation_requests(self): def _runner_stub(runner_type=DSparkCUDAGraphRunner): runner = runner_type.__new__(runner_type) - runner.clear = lambda: None + runner.clear = lambda **_kwargs: None return runner @@ -220,11 +220,11 @@ def test_decoder_selects_confidence_graph_subclass_only_when_enabled( decoder = SimpleNamespace( _dspark_confidence_enabled=confidence_enabled, _cuda_graph_runner_config=lambda: config, - model=SimpleNamespace(register_cuda_graph_pre_replay_hooks=registered.append), + _model=SimpleNamespace(register_cuda_graph_pre_replay_hooks=registered.append), ) runner = DecoderRunner._initialize_cuda_graph_runner(decoder) - runner.clear = lambda: None + runner.clear = lambda **_kwargs: None expected_type = DSparkCUDAGraphRunner if confidence_enabled else CUDAGraphRunner assert type(runner) is expected_type diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py index 54c258940dc1..50ca21edcc05 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py @@ -373,7 +373,7 @@ def test_model_engine_delegates_ragged_state_to_decoder_owner() -> None: inputs._dspark_exact_identity_words = tuple(range(8)) inputs.fit_ragged_verify_lens = lambda *args, **kwargs: (args, kwargs) worker = object() - runner.model = SimpleNamespace(spec_worker=worker) + runner._model = SimpleNamespace(spec_worker=worker) engine = object.__new__(PyTorchModelEngine) engine._cleanup_done = True engine._runner = runner @@ -440,7 +440,7 @@ def __init__(self, **kwargs): metadata_module.build_attention_metadata(model, **kwargs) assert "enable_ragged_verification" not in seen[0] - assert seen[0]["num_heads_per_kv"] == 4 + assert seen[0]["max_num_heads_per_kv"] == 4 def test_metadata_helper_enables_only_supported_ragged_metadata(): @@ -481,6 +481,7 @@ def __init__(self, **kwargs): **kwargs, spec_config=SimpleNamespace(enable_confidence_scheduling=enabled) ) decoder.mapping = kwargs["mapping"] + decoder._attention_runtime_features = kwargs["attention_runtime_features"] decoder._buffers = SimpleNamespace(cache_indirection=None) monkeypatch.setattr(decoder_metadata_module, "resolve_mamba_metadata_cls", lambda _model: None) kv = object() @@ -498,7 +499,7 @@ def __init__(self, **kwargs): @pytest.mark.parametrize("graph_type", [CUDAGraphRunner, DSparkCUDAGraphRunner]) def test_current_main_warmup_keeps_serving_timeout_only_for_capture(monkeypatch, graph_type): graph = graph_type.__new__(graph_type) - graph.clear = lambda: None + graph.clear = lambda **_kwargs: None graph._capture_allowed = False graph.padding_dummy_requests = {5: object()} if graph_type is DSparkCUDAGraphRunner: @@ -569,12 +570,16 @@ def test_confidence_without_an_sps_table_is_eager_only() -> None: graphs._get_graphs_to_capture([1, 4]) +@pytest.mark.parametrize("use_fused_gather", [False, True]) @pytest.mark.parametrize( ("trims", "prefix_window", "previous_windows"), [(True, 1, [3, 3]), (True, 1, [1, 2]), (False, None, [None, None])], ) def test_overlap_offsets_follow_the_packed_prefix( - trims: bool, prefix_window: int | None, previous_windows: list[int | None] + trims: bool, + prefix_window: int | None, + previous_windows: list[int | None], + use_fused_gather: bool, ) -> None: """Exercise the full input packer with CPU buffers and metadata-only leaves.""" decoder = object.__new__(DecoderInputPreparer) @@ -618,6 +623,27 @@ def test_overlap_offsets_follow_the_packed_prefix( decoder._mm_item_scheduler = None decoder._pinned_host_cache = {} decoder._pinned_host_active = {} + decoder._staged_previous_batch_indices = None + decoder._staged_previous_pos_indices = None + gather_calls = [] + + def gather(*args) -> bool: + # A CPU kernel leaf verifies the production caller's packed offsets. + new, draft, lens, slots, positions, rows, width, draft_width = args[:8] + ids, input_begin, drafts, draft_begin, pos, pos_begin, kv, kv_begin = args[8:] + selected_slots = slots[:rows].long() + gather_calls.append((input_begin, pos_begin, kv_begin)) + ids[input_begin : input_begin + rows * width].copy_( + new.transpose(0, 1)[selected_slots, :width].flatten() + ) + drafts[draft_begin : draft_begin + rows * draft_width].copy_( + draft[selected_slots, :draft_width].flatten() + ) + pos[pos_begin : pos_begin + rows * width].copy_(lens[positions[: rows * width].long()]) + kv[kv_begin : kv_begin + rows].copy_(lens[selected_slots] - width) + return True + + decoder._step_input_gather = SimpleNamespace(gather=gather) if use_fused_gather else None decoder._ragged_gather_indices = lambda slots, lengths: tuple( torch.tensor(indices, dtype=torch.long) for indices in ragged_gather_index_lists(slots, lengths) @@ -626,7 +652,6 @@ def test_overlap_offsets_follow_the_packed_prefix( for name in ( "input_ids_cuda", "draft_tokens_cuda", - "num_accepted_draft_tokens_cuda", "previous_pos_id_offsets_cuda", "previous_kv_lens_offsets_cuda", "previous_pos_indices_cuda", @@ -635,6 +660,7 @@ def test_overlap_offsets_follow_the_packed_prefix( "gather_ids_cuda", ): setattr(decoder._buffers, name, torch.full((32,), -99, dtype=torch.int32)) + decoder._gather_ids_buffer = decoder_module.DeviceCopyBuffer(decoder._buffers.gather_ids_cuda) windows = [prefix_window, *previous_windows] requests = [ SimpleNamespace( @@ -699,3 +725,87 @@ def test_overlap_offsets_follow_the_packed_prefix( assert metadata.kv_cache_params.num_extra_kv_tokens == 2 assert gather_ids.tolist() == list(range(sum(lengths))) assert draft_len == 3 + assert gather_calls == ( + [(prefix_len, prefix_len, 1)] if use_fused_gather and previous_lengths == [4, 4] else [] + ) + + +@pytest.mark.parametrize("trims", [False, True]) +def test_graph_overlap_offset_zeroing_preserves_ragged_bucket_padding(trims: bool) -> None: + decoder = object.__new__(DecoderInputPreparer) + decoder._config = SimpleNamespace( + max_num_tokens=16, max_beam_width=1, disable_overlap_scheduler=False + ) + decoder._dspark_trims_submitted_tokens = trims + decoder._staged_previous_batch_indices = None + decoder._staged_previous_pos_indices = None + decoder._step_input_gather = None + decoder.get_runtime_tokens_per_gen_step = lambda _length: 4 + decoder._ragged_gather_indices = lambda slots, lengths: tuple( + torch.tensor(indices, dtype=torch.long) + for indices in ragged_gather_index_lists(slots, lengths) + ) + decoder._buffers = SimpleNamespace( + **{ + name: torch.full((16,), -99, dtype=torch.int32) + for name in ( + "input_ids_cuda", + "draft_tokens_cuda", + "previous_batch_indices_cuda", + "previous_pos_indices_cuda", + "previous_pos_id_offsets_cuda", + "previous_kv_lens_offsets_cuda", + ) + } + ) + lengths = [1, 2] if trims else [4, 4] + host = SimpleNamespace( + lists=SimpleNamespace( + input_ids=[], + draft_tokens=[], + draft_lens=[], + position_ids=list(range(sum(lengths))), + previous_batch_indices=[0, 1], + previous_pos_indices=[0] * lengths[0] + [1] * lengths[1], + extend_requests=[object(), object()], + extend_dummy_requests=[], + generation_requests=[], + unwritten_input_id_rows=[], + ), + previous_batch_tokens_per_request=lengths if trims else None, + previous_batch_token_start=0, + ) + decoder._copy_token_ids_to_device( + host, + SimpleNamespace(is_attention_dp_phase_idle=False, num_context_requests=0), + torch.arange(8, dtype=torch.int32).reshape(4, 2, 1), + torch.tensor([1, 2], dtype=torch.int32), + torch.arange(6, dtype=torch.int32).reshape(2, 3), + 3, + attn_metadata=SimpleNamespace(is_cuda_graph=True), + enable_spec_decode=True, + ) + expected_tail = 0 if trims else -99 + assert decoder._buffers.previous_pos_id_offsets_cuda[sum(lengths) :].tolist() == [ + expected_tail + ] * (16 - sum(lengths)) + assert decoder._buffers.previous_kv_lens_offsets_cuda[2:].tolist() == [expected_tail] * 14 + + +@pytest.mark.parametrize("release_owners", [False, True]) +def test_graph_release_preserves_nccl_window_owner_choice(release_owners: bool) -> None: + calls = [] + graphs = object.__new__(DecoderCudaGraphs) + graphs.cuda_graph_runner = SimpleNamespace( + clear=lambda **kwargs: calls.append(("native", kwargs)) + ) + graphs.breakable_cuda_graph_runner = SimpleNamespace( + clear=lambda **kwargs: calls.append(("breakable", kwargs)) + ) + + graphs.release(release_nccl_window_owners=release_owners) + + assert calls == [ + ("native", {"release_nccl_window_owners": release_owners}), + ("breakable", {"release_nccl_window_owners": release_owners}), + ]