From 6d85142f92fdce7810f7133c5277a415aa0b8038 Mon Sep 17 00:00:00 2001 From: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> Date: Thu, 15 Oct 2026 11:45:01 -0700 Subject: [PATCH 1/7] [TRTLLM-14620][feat] Integrate reviewed DSpark graph runtime on current main Local unpublished integration candidate based on canonical03748030. Applies the published PR5 609be997 changes with native-off decoder behavior retained. Native validation remains pending. 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/runner.py | 920 +++++++++++++++++- .../_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 + .../executor/test_pytorch_model_engine.py | 7 + .../test_cuda_graph_dummy_identity.py | 94 ++ .../test_dspark_cuda_graph_runtime.py | 303 ++++++ .../test_dspark_model_engine_runtime.py | 629 ++++++++++++ .../hw_agnostic/test_dspark_select.py | 71 ++ .../test_device_window_sampling_layout.py | 112 +++ 14 files changed, 2858 insertions(+), 112 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 19ed990f33a1..f84f2c058026 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -28,8 +28,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 @@ -46,6 +46,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, EncoderDecoderRunner._prepare_capture_batch temporarily runs each # dummy generation request as a one-token context chunk to write its cross-KV @@ -58,6 +84,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 @@ -76,6 +141,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( @@ -430,7 +498,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: @@ -507,6 +576,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) @@ -568,28 +663,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 @@ -989,6 +1066,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: @@ -1015,20 +1111,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 @@ -1041,6 +1124,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 @@ -1065,12 +1149,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. @@ -1078,8 +1206,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) @@ -1104,8 +1233,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] @@ -1122,17 +1253,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( @@ -1169,20 +1326,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: @@ -1241,6 +1405,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, @@ -1264,11 +1434,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 @@ -1338,12 +1503,207 @@ def clear(self): 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 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 c9ed43c82ef9..9cd712e5ca85 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/metadata.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/metadata.py @@ -53,6 +53,7 @@ def build_attention_metadata( draft_kv_cache_manager: KVCacheManager | KVCacheManagerV2 | None = None, enable_context_mla_with_cached_kv: bool | None = None, 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 @@ -69,6 +70,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, @@ -82,4 +92,5 @@ def build_attention_metadata( cache_indirection=cache_indirection, num_heads_per_kv=num_heads_per_kv, sparse_metadata_params=sparse_metadata_params, + **metadata_kwargs, ) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index c74a0a8d9e7e..49ee7677e181 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -11,7 +11,7 @@ import os from collections.abc import Callable, Iterator, Sequence from contextlib import contextmanager -from typing import Any +from typing import Any, List, Optional, Tuple import torch @@ -38,6 +38,7 @@ from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( CUDAGraphRunner, CUDAGraphRunnerConfig, + DSparkCUDAGraphRunner, KeyType, get_mrope_dummy_seq_slot, ) @@ -47,7 +48,11 @@ BaseMambaCacheManager, MambaHybridCacheManager, ) -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 ( BaseResourceManager, KVCacheManager, @@ -71,6 +76,7 @@ ) 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, resolve_draft_len from tensorrt_llm._torch.utils import ( @@ -485,6 +491,14 @@ def _allocate_decoder_buffers(self) -> None: self.previous_batch_indices_cuda = torch.empty( (self._config.max_num_tokens,), dtype=torch.int, device="cuda" ) + self.ragged_verify_lens_cuda = torch.empty( + (self._config.max_batch_size + 1,), dtype=torch.int, device="cuda" + ) + self.ragged_qo_indptr_cuda = torch.empty( + (self._config.max_batch_size + 2,), dtype=torch.int, device="cuda" + ) + self._pinned_host_cache = {} + self._pinned_host_active = {} self.input_ids_cuda = torch.empty( (self._config.max_num_tokens,), dtype=torch.int, device="cuda" ) @@ -512,6 +526,37 @@ def _allocate_decoder_buffers(self) -> None: def _init_decoder_state(self) -> None: """Initialize the mutable state read only by engine decoder execution.""" + 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 self._eager_workspace_reclaimer: EagerWorkspaceReclaimer | None = None # Token shape of the last autotuner warmup that ran a forward pass. self._completed_autotuner_warmup_num_tokens: int | None = None @@ -542,7 +587,8 @@ def _init_decoder_state(self) -> None: self._dynamic_draft_len_mapping = self._compute_dynamic_draft_len_mapping() 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) @@ -773,7 +819,7 @@ def _warmup_scheduled(self, resource_manager: ResourceManager, kv_cache_manager) # The lifetime of model engine and kv cache manager can be different. # Reset the global cuda graph dummy requests in warmup. - self.cuda_graph_runner.padding_dummy_requests = {} + self.cuda_graph_runner.reset_padding_dummies() if self.mapping.cp_size > 1: cp_type = self.mapping.cp_config.get("cp_type", None) @@ -881,6 +927,7 @@ def _warmup_scheduled(self, resource_manager: ResourceManager, kv_cache_manager) # 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) + self._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") @@ -1956,7 +2003,9 @@ def _compute_dynamic_draft_len_mapping(self) -> dict[int, int] | None: mapping[graph_bs] = draft_len return mapping - 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: @@ -1996,6 +2045,33 @@ def _get_graphs_to_capture(self, cuda_graph_batch_sizes: list[int]) -> list[tupl ) return graphs + if self._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._dspark_trims_submitted_tokens: + return native_graphs + from tensorrt_llm._torch.speculative.dspark_planner import ExactSpsCostTable + + exact_table = self._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 _prepare_tp_inputs: # logical K for linear-tree modes, total tree tokens for tree decoding. @@ -2025,7 +2101,7 @@ def _warmup_and_capture_cuda_graphs(self, resource_manager: ResourceManager) -> self._run_cuda_graph_warmup(resource_manager) finally: self.cuda_graph_runner.is_warmup_only = False - self.cuda_graph_runner.padding_dummy_requests = {} + self.cuda_graph_runner.reset_padding_dummies() # Captured graphs keep the MoE communication timeouts that were # current at capture, and serving replays them, so capture with the # serving timeouts. The warmup-only pass above keeps the warmup ones. @@ -2162,7 +2238,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 @@ -2175,6 +2253,8 @@ def _run_capture_pass( force_non_greedy=force_non_greedy, ) with self._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 @@ -2211,6 +2291,8 @@ def _run_capture_pass( finally: self._force_lora_graph_for_capture = None self.cuda_graph_runner.set_capture_sample_type(None) + if self._dspark_confidence_enabled: + self.cuda_graph_runner.agreed_ragged_bucket = None if self._lora.cuda_graph_manager is None: lora_graph_cases = [False] @@ -2790,6 +2872,11 @@ def _add_dummy_request_resources( """Allocate resources dummy requests need beyond the KV caches.""" return True + @property + def _ragged_verification_enabled(self) -> bool: + """Select persistent ragged storage before attention metadata construction.""" + return self._dspark_confidence_enabled + def _set_up_attn_metadata( self, kv_cache_manager: KVCacheManager | KVCacheManagerV2, @@ -2815,6 +2902,7 @@ def _set_up_attn_metadata( 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( @@ -3130,6 +3218,662 @@ def _record_cached_kv_tokens_per_req( num_cached_tokens_per_seq ) - sum(counts) + @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._dspark_confidence_enabled or not self._dspark_trims_submitted_tokens: + 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", + "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 = [] + 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._compute_compressed_mask( + metadata.new_comp_kv_lens_cuda, + metadata.cu_new_comp_kv_cuda, + metadata.compressed_mask_cuda, + num_generations, + total_compressed_tokens, + compress_ratios, + ) + torch.cuda.synchronize() + finally: + 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 _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) + + 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 _get_spec_worker(self): + """Return the one-model speculative worker when it is initialized.""" + return getattr(self.model, "spec_worker", None) + + def fit_ragged_verify_lens( + self, + generation_requests, + verify_lens: List[int], + exact_shape: Tuple[int, int, int], + peer_stats: Optional[List[List[int]]] = None, + exact_zero_real: bool = False, + ) -> Optional[int]: + """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 _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.ragged_verify_lens_cuda + qo_buf = self.ragged_qo_indptr_cuda + split_lens = lens_buf[:padded_bs].clone() + split_qo = qo_buf[: padded_bs + 1].to(torch.long) + past_seen = self.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.previous_kv_lens_offsets_cuda[:padded_bs] -= window_delta.to( + self.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.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.input_ids_cuda[:bucket] = ( + new_tokens_device.transpose(0, 1)[slots_tok, offset] + .flatten() + .to(self.input_ids_cuda.dtype) + ) + self.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.previous_pos_indices_cuda[:real_tokens] = slots_tok[:real_tokens].to( + self.previous_pos_indices_cuda.dtype + ) + self.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.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.draft_tokens_cuda.dtype) + ) + + apply_device_layout(result.verify_lens, req_idx, result.kv_correction) + return True + + @staticmethod + def _ragged_token_lens(generation_requests) -> Optional[List[int]]: + """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.ragged_verify_lens_cuda[:n] + lens_view.copy_( + self._pinned_host("ragged_verify_lens", token_lens, torch.int32), non_blocking=True + ) + indptr_view = self.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) + def _can_use_steady_gen_fast_prepare( self, scheduled_requests: ScheduledRequests, @@ -3673,6 +4417,9 @@ def _helix_pack_extend(request, group: int) -> int: # attention-DP dummies) whose input ids are not appended to the host # ``input_ids`` list; zero-filled after the input-id copies below. unwritten_input_id_rows = [] + previous_batch_token_start = None + previous_batch_tokens_per_request = [] if self._dspark_trims_submitted_tokens else None + extend_tokens_emitted = 0 runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step(runtime_draft_len) runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 for request in extend_requests: @@ -3692,6 +4439,11 @@ def _helix_pack_extend(request, group: int) -> int: 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): @@ -3699,10 +4451,12 @@ def _helix_pack_extend(request, group: int) -> int: 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 + if self._dspark_trims_submitted_tokens: + submitted_draft_tokens = submitted_draft_tokens[:num_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 @@ -3737,6 +4491,9 @@ def _helix_pack_extend(request, group: int) -> int: num_cached_tokens_per_seq.append( past_seen_token_num - request.py_num_compressed_tokens ) + if self._dspark_trims_submitted_tokens: + request.cached_tokens = past_seen_token_num + extend_tokens_emitted += 1 + num_draft_tokens if _has_cp_helix: # Verify group [base, base+group) in GLOBAL positions. # On a helix gen worker the request's token list is the @@ -3760,30 +4517,36 @@ def _helix_pack_extend(request, group: int) -> int: 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 + ) + sequence_lengths.append(request_tokens_per_step) num_accepted_draft_tokens.append(request.py_num_accepted_draft_tokens) 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) + 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 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 @@ -3792,7 +4555,7 @@ def _helix_pack_extend(request, group: int) -> int: # 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) @@ -3802,7 +4565,7 @@ def _helix_pack_extend(request, group: int) -> int: 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) @@ -4056,6 +4819,9 @@ def _helix_pack_extend(request, group: int) -> int: mrope_delta_read_seq_slots.clear() 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(): previous_batch_indices_host = torch.tensor( @@ -4099,18 +4865,43 @@ def previous_seq_slots_device(): if previous_batch_len > 0: previous_slots = previous_seq_slots_device() + is_ragged_gen = previous_batch_tokens_per_request is not None and any( + tokens != runtime_tokens_per_gen_step + for tokens in previous_batch_tokens_per_request + ) # previous input ids - previous_batch_tokens = previous_batch_len * runtime_tokens_per_gen_step - new_tokens = new_tokens_device.transpose(0, 1)[ - previous_slots, :runtime_tokens_per_gen_step - ].flatten() + previous_batch_tokens = len(previous_pos_indices) + if is_ragged_gen: + token_rows, token_cols = self._ragged_gather_indices( + previous_batch_indices, 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.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 previous_batch_tokens_per_request], + ) + self.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.draft_tokens_cuda[ num_draft_tokens : num_draft_tokens + previous_batch_draft_tokens ].copy_( @@ -4120,7 +4911,18 @@ def previous_seq_slots_device(): 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( + 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] previous_pos_indices_host = torch.tensor( previous_pos_indices, dtype=torch.int, pin_memory=prefer_pinned() ) @@ -4149,10 +4951,14 @@ def previous_seq_slots_device(): # Already set to '0' during initialization. num_extend_reqeust_wo_dummy = len(extend_requests) - len(extend_dummy_requests) - self.previous_pos_id_offsets_cuda[ - (num_extend_reqeust_wo_dummy - previous_batch_len) - * runtime_tokens_per_gen_step : num_extend_reqeust_wo_dummy + pos_offsets_start = ( + previous_batch_token_start + if previous_batch_token_start is not None + else (num_extend_reqeust_wo_dummy - previous_batch_len) * runtime_tokens_per_gen_step + ) + self.previous_pos_id_offsets_cuda[ + pos_offsets_start : pos_offsets_start + previous_batch_tokens ].copy_( new_tokens_lens_device[self.previous_pos_indices_cuda[0:previous_batch_tokens]], non_blocking=True, @@ -4160,7 +4966,7 @@ def previous_seq_slots_device(): self.previous_kv_lens_offsets_cuda[ num_extend_reqeust_wo_dummy - previous_batch_len : 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() @@ -4254,7 +5060,21 @@ def previous_seq_slots_device(): helix_owned_new_tokens=(helix_owned_new_tokens if helix_spec_active else None), ) - 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=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(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( @@ -4484,6 +5304,10 @@ def previous_seq_slots_device(): ] if context_prompt_lookahead is not None: spec_metadata.populate_context_prompt_lookahead(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() @@ -4528,6 +5352,13 @@ def previous_seq_slots_device(): ) if not is_dummy: + self.previous_request_ids = 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 # Record the steady-state generation cache when this pass handled # purely non-dummy generation requests that all carried a previous # overlap-scheduler tensor (previous_batch_len == _n_gen implies @@ -4890,6 +5721,13 @@ def _forward_decoder( self._prepare_inputs_event = torch.cuda.Event() self._prepare_inputs_event.record() + if ( + can_run_graph + and not is_dummy + and getattr(self, "_dspark_device_budget", None) is not None + ): + self._apply_device_window_prologue(inputs, new_tensors_device) + breakable_runner = self.breakable_cuda_graph_runner with with_shared_pool(self.cuda_graph_runner.get_graph_pool()): @@ -5145,6 +5983,16 @@ def _execute_logit_post_processors(self, scheduled_requests: ScheduledRequests, 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 diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 221f9145b17c..f064391d0037 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -804,6 +804,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 f1273926d683..3e8690d8859e 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) @@ -558,8 +558,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 @@ -567,7 +568,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 8237c09ee8b7..ff345a77e7a4 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -1584,6 +1584,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 90f9bb55eb51..c77dff8414ac 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -89,6 +89,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_hf_quant_config.py diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index a2ed4a1b20eb..b76cd5c3b1dc 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -828,6 +828,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] @@ -856,6 +857,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 @@ -880,6 +882,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 @@ -927,6 +930,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] @@ -945,6 +949,7 @@ 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._dspark_confidence_enabled = False runner.config = SimpleNamespace( enable_attention_dp=False, use_mrope=False, @@ -985,6 +990,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] @@ -1027,6 +1033,7 @@ def test_lora_graph_variant_selection(self) -> None: def test_graph_lookup_forwards_promoted_context_ids(self) -> None: runner = Mock() runner.enabled = True + 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..0416bb58956b --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph_runtime.py @@ -0,0 +1,303 @@ +# 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, + is_first_draft=False, + 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, + is_first_draft=False, + 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..7e1b0ffde6a7 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py @@ -0,0 +1,629 @@ +# 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.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 runner as decoder_module +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__(DecoderRunner) + 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 = DecoderRunner._pinned_host(engine, "rows", [1, 2], torch.long) + second = DecoderRunner._pinned_host(engine, "rows", [3, 4], torch.long) + third = DecoderRunner._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 = DecoderRunner.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"): + DecoderRunner.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 = DecoderRunner.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 = DecoderRunner.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}, + _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(*args): + mask_calls.append((torch.is_inference_mode_enabled(), args[3:])) + + metadata = _warmup_metadata(record_positions, record_mask) + engine = SimpleNamespace( + _dspark_confidence_enabled=True, + _dspark_trims_submitted_tokens=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)) + + DecoderRunner._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: 42, 4: 14}, [1, 4]))] + 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( + _dspark_confidence_enabled=True, + _dspark_trims_submitted_tokens=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"): + DecoderRunner._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() + + +def test_ragged_compressor_warmup_requires_trimmed_tokens(): + engine = SimpleNamespace( + _dspark_confidence_enabled=True, + _dspark_trims_submitted_tokens=False, + ) + + DecoderRunner._warmup_dspark_ragged_compressor_metadata(engine) + + +def test_ragged_compressor_warmup_rejects_missing_metadata_contract(): + engine = SimpleNamespace( + _dspark_confidence_enabled=True, + _dspark_trims_submitted_tokens=True, + attn_metadata=SimpleNamespace(), + ) + + with pytest.raises(RuntimeError, match="missing .*max_draft_tokens"): + DecoderRunner._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.ragged_verify_lens_cuda = torch.empty(9, dtype=torch.int32) + runner.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: + runner = _engine(_runner(), []) + runner._dspark_sps_cost_table = object() + runner._dspark_exact_candidate_cells = ((4, 12),) + runner._dspark_exact_identity_words = tuple(range(8)) + runner.fit_ragged_verify_lens = lambda *args, **kwargs: (args, kwargs) + worker = object() + runner._get_spec_worker = lambda: 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__(DecoderRunner) + decoder.attn_metadata = None + decoder._dspark_confidence_enabled = enabled + decoder.model = SimpleNamespace(model_config=model) + decoder._config = SimpleNamespace(**kwargs) + decoder.mapping = kwargs["mapping"] + decoder.cache_indirection_attention = None + monkeypatch.setattr(decoder_module, "resolve_mamba_metadata_cls", lambda _model: None) + kv = object() + + result = decoder._set_up_attn_metadata(kv) + if enabled: + result.ragged_verify_lens = [1, 2] + result._validate_ragged_verification() + + assert seen == [enabled] + assert decoder._set_up_attn_metadata(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(decoder_module, "moe_comm_serving_timeouts", serving_timeouts) + decoder = SimpleNamespace(cuda_graph_runner=graph, _run_cuda_graph_warmup=warmup) + DecoderRunner._warmup_and_capture_cuda_graphs(decoder, 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__(DecoderRunner) + 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__(DecoderRunner) + 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) + ) + + decoder._capture_generation_cuda_graphs(None) + with pytest.raises(RuntimeError, match="authenticated exact SPS table"): + decoder._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__(DecoderRunner) + mode = SimpleNamespace( + support_overlap_scheduler=lambda: True, extend_ctx=lambda _backend: False + ) + spec_config = SimpleNamespace(spec_dec_mode=mode, is_linear_tree=True) + 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, 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.use_beam_search = False + decoder.iter_states = {} + decoder._log_cached_kv_tokens_per_req = False + 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) + ) + 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, 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) + + inputs, gather_ids, draft_len = decoder._prepare_tp_inputs( + scheduled, + None, + metadata, + new_tensors_device=state, + enable_spec_decode=True, + runtime_draft_len=3, + is_dummy=False, + ) + + 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.previous_pos_id_offsets_cuda[: sum(lengths)].tolist() == expected + assert decoder.previous_pos_id_offsets_cuda[sum(lengths) :].count_nonzero().item() == 0 + assert decoder.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 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/test_device_window_sampling_layout.py b/tests/unittest/_torch/speculative/test_device_window_sampling_layout.py new file mode 100644 index 000000000000..fe29ae5c95cb --- /dev/null +++ b/tests/unittest/_torch/speculative/test_device_window_sampling_layout.py @@ -0,0 +1,112 @@ +# 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 + + +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 cd2344b98972fae9eccfb24d39e1cdc0947adc6b Mon Sep 17 00:00:00 2001 From: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> Date: Thu, 15 Oct 2026 11:48:03 -0700 Subject: [PATCH 2/7] [TRTLLM-14620][feat] Reconcile pending DSpark executor activation Local unpublished PR6 integration checkpoint. Preserves current-main refactors and published PR5 graph subclass, selector rename, worker guard, packed prefixes and constructor-time DSA metadata. Matching-build golden/schema generation and runtime validation remain pending. Signed-off-by: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> --- examples/models/core/deepseek_v4/README.md | 205 ++++++ .../custom_ops/dspark_attention_custom_op.py | 11 +- .../dspark_rmsnorm_rope_custom_op.py | 2 +- .../blackwell/dspark/attention.py | 7 +- tensorrt_llm/_torch/models/modeling_dspark.py | 16 +- .../_torch/models/modeling_speculative.py | 13 +- tensorrt_llm/_torch/pyexecutor/_util.py | 13 + .../_torch/pyexecutor/cuda_graph_runner.py | 5 + .../pyexecutor/dspark_mixed_host_capacity.py | 181 +++++ .../pyexecutor/dspark_mixed_host_snapshot.py | 168 +++++ .../_torch/pyexecutor/dspark_mixed_runtime.py | 474 +++++++++++++ .../engine/runners/decoder/runner.py | 314 +++++++-- .../_torch/pyexecutor/model_engine.py | 7 + tensorrt_llm/_torch/pyexecutor/py_executor.py | 624 +++++++++++++++++- tensorrt_llm/_torch/speculative/dspark.py | 62 +- .../speculative/dspark_mixed_evidence.py | 514 +++++++++++++++ .../_torch/speculative/dspark_planner.py | 82 ++- .../_torch/speculative/dspark_schedule.py | 279 ++++++++ .../_torch/speculative/dspark_select.py | 525 ++++++++++++++- .../_torch/speculative/dspark_verify.py | 34 +- tensorrt_llm/llmapi/llm_args.py | 327 ++++++++- .../usage/llm_args_golden_manifest.json | 15 + .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../_torch/executor/engine/test_metadata.py | 44 ++ .../test_kv_cache_v2_capacity_only.py | 22 + .../test_cuda_graph_dummy_identity.py | 4 + .../test_dspark_admission_resolution.py | 121 ++++ .../hw_agnostic/test_dspark_artifact_reuse.py | 197 ++++++ .../test_dspark_cuda_graph_runtime.py | 82 +++ .../test_dspark_dummy_slots_cpu.py | 75 +++ ...test_dspark_embedded_confidence_loading.py | 102 +++ .../test_dspark_executor_policy.py | 326 +++++++++ .../test_dspark_fused_attention_route.py | 328 +++++++++ .../hw_agnostic/test_dspark_mixed_evidence.py | 458 +++++++++++++ .../hw_agnostic/test_dspark_mixed_runtime.py | 350 ++++++++++ .../test_dspark_sampling_admission.py | 145 ++++ .../hw_agnostic/test_dspark_schedule_fused.py | 543 +++++++++++++++ .../hw_agnostic/test_dspark_select.py | 30 + .../hw_agnostic/test_dspark_sts.py | 10 + .../hw_agnostic/test_dspark_worker.py | 19 +- .../test_dspark_worker_stamp_wrap_cpu.py | 59 ++ .../test_dspark_cute_dsl_attention.py | 109 ++- .../test_dspark_workload_admission_config.py | 140 ++++ tests/unittest/llmapi/test_llm_args.py | 192 +++++- 44 files changed, 7050 insertions(+), 185 deletions(-) create mode 100644 tensorrt_llm/_torch/pyexecutor/dspark_mixed_host_capacity.py create mode 100644 tensorrt_llm/_torch/pyexecutor/dspark_mixed_host_snapshot.py create mode 100644 tensorrt_llm/_torch/pyexecutor/dspark_mixed_runtime.py create mode 100644 tensorrt_llm/_torch/speculative/dspark_mixed_evidence.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_admission_resolution.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_artifact_reuse.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_dummy_slots_cpu.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_embedded_confidence_loading.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_executor_policy.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_fused_attention_route.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_evidence.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_runtime.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sampling_admission.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_schedule_fused.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_stamp_wrap_cpu.py create mode 100644 tests/unittest/llmapi/test_dspark_workload_admission_config.py diff --git a/examples/models/core/deepseek_v4/README.md b/examples/models/core/deepseek_v4/README.md index 5672ef03a73e..ddac8fa83a45 100644 --- a/examples/models/core/deepseek_v4/README.md +++ b/examples/models/core/deepseek_v4/README.md @@ -472,6 +472,211 @@ MXFP4 routed MoE expert weights, TensorRT LLM automatically applies the routed-e configuration. +## DSpark confidence verification (prototype) + +Confidence-guided verification is **default-off and experimental**. It proposes the full physical +draft block `K`, then may select a smaller, measured verifier-token budget. It does not reduce the +physical drafter length. This draft integration is not a production deployment recipe: strict +receipt generation and performance qualification are external, and the composed source still needs native +build, GPU/full-model, accuracy and performance validation. Optional mixed scheduling has the +separate evidence and route requirements below; no composed-tip GPU qualification is implied. + +`K` is the physical draft length, `G` is the rank-local request count for a captured graph, and +`V` is its verifier-token budget. In the cost table, `V=0` denotes the native full-width path; +positive `V` values denote measured compact layouts. The table must cover the configured physical +`K` and exact graph ladder. The scheduler never invents a missing cell. + +### Requested configuration versus actual execution + +Set `DSparkDecodingConfig.enable_confidence_scheduling=True` to opt in. By default, structurally +valid existing costs and calibration can be reused even when their recorded source or GPU-node +metadata differ from the current run. A warning makes clear that predictions may be stale and +that positive performance economics have not been established. Reuse does not change target +verification or the acceptance rule. + +Set `validate_confidence_artifact_freshness=True` to require the previous strict live-fingerprint +and positive workload-admission checks before constructing engines and workers: + +| Startup result | Resolved execution | +| --- | --- | +| Explicit opt-in with default artifact reuse | Confidence stays enabled; measured costs must still pass schema, finiteness, physical-`K`, exact graph and layout checks. No admission receipt is created or asserted. | +| Strict validation with missing, unreadable, stale or non-positive admission evidence | Both confidence flags become false and all confidence artifact paths/digests are cleared; ordinary feature-off static `K` runs. A warning explains the fallback. | +| Physical `K=1` | Feature-off static `K1`; there is no smaller verifier tier. Confidence inputs are ignored and cleared. | +| Strict validation with a matching receipt and positive aggregate value | Confidence remains requested; the existing environment, graph, checkpoint and runtime checks must still pass. | + +An **enabled-native** policy decision is different from startup +feature-off fallback: it retains the enabled control path while verifying the native layout. +An enabled compact decision uses a positive measured budget. Optional fused-scheduler failure +falls back to the tensor scheduler for the affected graph size; that is a third, separate fallback. + +Before reporting an "enabled" experiment, inspect the resolved +`TorchLlmArgs.speculative_config.enable_confidence_scheduling` and record both the requested and +resolved configuration. Require observed native/compact routing and completed outputs, rather +than inferring execution from an input YAML flag or a successful server start. Do not count +startup feature-off fallback as an enabled-native run. With confidence disabled, omit all +confidence-only fields; explicitly supplying them is a configuration error. + +Confidence scheduling requires padded CUDA graphs and PP=1. Guided decoding, LoRA and incompatible +sparse top-k settings remain unsupported. Per-request logits post-processors are rejected before +activation. Advanced rejection draft sampling is also unsupported: its sampled proposal prefix +differs from the greedy prefix that produced the confidence scores. Greedy requests (including +`top_k=1`) and requests with rejection draft sampling disabled retain their existing behavior. +Confidence-disabled execution retains its ordinary sampling and post-processing paths. + +### Reusing costs and calibration + +For default reuse, supply `confidence_sps_table_path` and, when using a calibrated head, +`confidence_sts_path`. Costs must be finite positive schema-v2 measurements with native and compact +cells for every configured graph size; STS temperatures must be a finite positive vector with +one value per physical draft position. Invalid artifacts raise errors rather than being silently +timed as enabled static execution. Their original metadata is retained; the runner records its +observed runtime separately instead of rewriting the artifact's provenance. + +For example, the confidence portion of a padded-graph configuration is: + +```yaml +speculative_config: + decoding_type: DSpark + max_draft_len: 3 + enable_confidence_scheduling: true + confidence_sps_table_path: /your/artifacts/k3-costs.json + confidence_sts_path: /your/artifacts/k3-sts.json +``` + +This is a configuration example, not an exact-tip GPU-validated reproduction recipe. Keep the +same checkpoint, physical `K`, inputs, parallelism, graph ladder and other runtime options when +comparing it with static DSpark (`enable_confidence_scheduling: false`, without confidence-only +fields). A fair measured comparison keeps source, configuration, inputs and output work matched, +uses independent model lifetimes, checks actual routing and completed outputs, and records actual +timings. Reused predictions do not invalidate that measurement: retain their original provenance +and label freshness unchecked, without promising the earlier gain. Fresh profiling is optional +and can improve prediction quality. Changing batch size requires a graph ladder covered by the +supplied table, not a separate per-batch-size package. + +Both modes retain the table's immediate predicted-goodput margin (normally 1%). If the table +supplies an `iteration_drain_guard`, its prices and final-drain policy are unchanged. Without that +optional metadata, default reuse selects by immediate goodput alone; strict validation retains +native verification. Neither behavior certifies an end-to-end gain. + +### Optional freshness and workload-admission validation + +With `validate_confidence_artifact_freshness: true`, provide all four prototype inputs: + +- `confidence_sps_table_path`: schema-v2 measured costs with exact physical `K`, rank-local graph + sizes `G`, native `V=0` cells and supported positive verifier budgets. +- `confidence_sps_live_fingerprint_path`: an independently supplied fingerprint for that runtime + and graph ladder, not a copy relabeled from an older cost artifact. +- `confidence_admission_receipt_path`: the schema-v1 workload receipt. +- `confidence_admission_receipt_sha256`: the expected lowercase SHA256 of the **receipt file + bytes**, obtained from a trusted, separately reviewed deployment package. + +There is no supported public cost-profile/calibration/receipt-generation command in this source +change. Self-generation requires the separate profiling tooling and its validated instructions; +the fields above are consumers, not a cost collector. Do not manufacture costs or an admitting +receipt from unit-test fixtures. A strict receipt's producer must preserve the measured cost, +workload, calibration, selector-replay and runtime evidence and review their applicability to the +actual deployment. A source, calibration, physical-`K` or workload change requires reassessment; +historical scores or timing results cannot be relabeled as new-source qualification. + +The receipt envelope contains exactly `schema_version: 1`, `admission` and `admission_sha256`. +The last field hashes the canonical admission-body JSON +(`json.dumps(body, sort_keys=True, separators=(",", ":")).encode()`). +The body contains: + +- `physical_k`, `source_head`, `source_diff_sha256`, `runtime_snapshot`; +- `sps_cost_table_sha256`, `live_engine_fingerprint_sha256`, `selector_identity_sha256`; +- `workload_identity_sha256`, `calibration_result_sha256`, `selector_replay_sha256`; +- `policy_steps`, `compact_choices`, `gross_compact_value_ms_lower_bound`, + `fixed_confidence_path_overhead_ms_upper_bound`, `safety_margin_ms`, `admitted`. + +The selector identity must match `EXACT_SPS_SELECTOR_IDENTITY_SHA256` in +`tensorrt_llm/_torch/speculative/dspark_planner.py`. Replay the existing marginal selector without +allocating fixed overhead back to candidates through a `policy_steps / compact_choices` ratio. +Admission recomputes: + +```text +net_value_lower_bound_ms = + gross_compact_value_ms_lower_bound + - policy_steps * fixed_confidence_path_overhead_ms_upper_bound + - safety_margin_ms + +admitted = compact_choices > 0 and net_value_lower_bound_ms > 0 +``` + +Native decisions still pay the enabled confidence-path overhead, so **all** policy steps count. +Strict admission does not replace per-decision selection: exact measured tiers and the existing +iteration/drain guard still determine whether a compact choice is allowed. + +Hashes provide integrity and cross-artifact consistency checks, **not signatures or proof of +origin**. The runtime checks supplied receipt, cost and fingerprint values; it does not independently +derive the running checkpoint/workload identity, re-execute the calibration or selector replay, +or verify their referenced result files. The external producer/deployment package remains the +trust boundary. Hashing an arbitrary receipt and passing that same digest does not establish +trust or qualify an enabled experiment. + +Contract regressions are in `test_dspark_workload_admission.py`, +`test_dspark_admission_resolution.py` (both under +`tests/unittest/_torch/speculative/hw_agnostic/`) and +`tests/unittest/llmapi/test_dspark_workload_admission_config.py`. These fixtures are synthetic +validator tests, not deployment evidence. `test_dspark_artifact_reuse.py` covers optional +freshness, unchanged structural guards and the absent-drain behavior. + +### Optional mixed context/generation scheduling + +`confidence_mixed_evidence` is an optional nested object. Its default is `None`, which leaves +mixed steps at native verification width and does not change enabled decode-only scheduling. +Mixed costs are not decode `T(G,V)` costs: context chunks and their positions, generation history, +bootstrap requests, ordered ranks and the shared body capacity all affect the mixed cost key. + +The object contains `costs_path`, `costs_sha256`, `admission_path`, `admission_sha256`, +`protocol_sha256`, `live_fingerprint_path` and `live_fingerprint_sha256`. Cost, receipt and live +fingerprint digests pin exact file bytes; the protocol digest identifies the measured protocol. +These pins come from a trusted deployment package. The live fingerprint must be observed independently +for the current source, runtime, model, GPU class, topology and selector, not copied from the cost +table. The cost file uses `dspark-mixed-breakable-costs-v1`; its separate receipt uses +`dspark-mixed-breakable-admission-v1`. The reuse setting applies to decode schema-v2 inputs only; +it does not relax this separate mixed-execution contract. + +The supported mixed route is embedded DSpark with physical `K=3`, TP8, PP1, CP1, attention DP, +greedy sampling and captured BREAKABLE body buckets `[96, 128, 192, 256, 384, 512]`. Standalone +draft models, unsupported routes and uncovered mixed geometry retain native mixed verification. +Missing or invalid mixed evidence does not disable an already enabled decode policy. +First context chunks also remain native because KV preparation may still resolve a prefix-cache +hit and change their geometry. Later chunks are eligible only when their full-width native rows +fit the configured token budget, so resource preparation does not trim a quoted context shape. + +When a peer rank prefills, all ranks use the existing shared BREAKABLE padding path rather than +replaying a rank-local full decode graph. Only authenticated generation owners may receive shorter +verification prefixes; context rows, bootstrap requests and graph-padding requests retain their +required rows. The runner publishes those lengths as host-authoritative before attention metadata +preparation, without changing the attention backend, sampler or acceptance rule. Stale or unready +confidence owners retain full physical `K`. + +Selection requires an exact measured geometry/action cell and retains the minimum 1% predicted +goodput gate. The current pricing field is `rank0_forward_pricing_proxy`, not a measurement of the +whole distributed step. Workload admission must charge the opt-in path's overhead on **all** policy +steps, including all-decode and native choices: the optional path adds CPU curve preparation, +producer-stamp transfer and a fixed trailer to the existing policy allgather. Policy equivalence +alone is not evidence of zero overhead or preserved throughput. + +There is no public mixed receipt-generation command in this change. Supplying valid-shaped JSON +or hashing test fixtures does not qualify a deployment. A separately reviewed producer must supply +matching measured costs, runtime/accuracy results, selector replay and positive workload economics. +Exact-tip GPU/full-model validation and production mixed evidence remain required before making a +deployment or performance claim. The CPU contracts are in `test_dspark_mixed_evidence.py` and +`test_dspark_mixed_runtime.py` under `tests/unittest/_torch/speculative/hw_agnostic/`. + +### Drafter fusion scope + +The BF16 CuTe attention and RMSNorm–RoPE path is eligible for physical `K=1..8` under the existing +SM100/SM103, geometry and layout guards; unsupported inputs retain the reference path. This +dispatch is independent of the confidence flag, so it also affects static DSpark. The eight-row +storage bound is unchanged; no arbitrary-`K` kernel extension is claimed. + +The source-matched extension passed 98 GPU unit cases on B300/SM103. That is not full-model or +performance qualification of this composed integration. SM100 and exact-integration GPU/model +validation remain pending; keep this scope experimental until those gates and review complete. + ## Notes and Troubleshooting - `DeepSeek-V4 on Hopper requires kv_cache_config.dtype='fp8_ds_mla'`: remove an explicit BF16/FP8 diff --git a/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py index d414a99f60a8..8feab7b891bb 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py @@ -5,7 +5,7 @@ The op runs the tcgen05 MMA kernel (:class:`DSparkAttention`) on the supported DSV4 DSpark -geometry: 128 heads, head_dim 512, draft block 5 or 6, and a 128-row rolling +geometry: 128 heads, head_dim 512, draft block 1 through 8, and a 128-row rolling window. Other shapes fall back to the pure-PyTorch reference path in ``models/modeling_dspark.py``. """ @@ -32,7 +32,8 @@ _DSV4_DSPARK_HEAD_DIM = 512 _DSV4_DSPARK_WINDOW_SIZE = 128 _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE = 8 -_DSV4_DSPARK_BLOCK_SIZES = (5, 6) +# Physical K is bounded by the eight-row draft-buffer ABI, not the MMA tile. +_DSV4_DSPARK_BLOCK_SIZES = (1, 2, 3, 4, 5, 6, 7, 8) _DSV4_DSPARK_ROPE_DIM = 64 @@ -92,7 +93,9 @@ def is_fused_dsv4_dspark_attention_supported( f"expected q/kv_cache ranks 4/3, got {q.ndim}/{kv_cache.ndim}", "tensor_ranks" ) if q.shape[1] not in _DSV4_DSPARK_BLOCK_SIZES: - return _log_unsupported(f"draft block size must be 5 or 6, got {q.shape[1]}", "block_size") + return _log_unsupported( + f"draft block size must be 1 through 8, got {q.shape[1]}", "block_size" + ) if q.shape[2:] != (_DSV4_DSPARK_NUM_HEADS, _DSV4_DSPARK_HEAD_DIM): return _log_unsupported( "q must have 128 heads and head_dim 512; " @@ -346,7 +349,7 @@ def fused_dsv4_dspark_attention( ): raise ValueError( "fused_dsv4_dspark_attention requires contiguous supported DSV4 DSpark tensors " - "([B, 5|6, 128, 512] queries, [B, 8, 512] draft blocks, INT32 indices, " + "([B, K, 128, 512] queries with 1 <= K <= 8, [B, 8, 512] draft blocks, INT32 indices, " "and a 128-row BF16 window) on SM100 or SM103" ) diff --git a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py index 56bc9f62d707..f2e0b741b2b5 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py @@ -20,7 +20,7 @@ _DSV4_DSPARK_ROPE_DIM = 64 _DSV4_DSPARK_WINDOW_SIZE = 128 _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE = 8 -_DSV4_DSPARK_BLOCK_SIZES = (5, 6) +_DSV4_DSPARK_BLOCK_SIZES = (1, 2, 3, 4, 5, 6, 7, 8) _DSV4_DSPARK_ARCH_BY_SM = { 100: "sm_100", 103: "sm_103", diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark/attention.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark/attention.py index 458d2877f12d..098458969918 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark/attention.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark/attention.py @@ -195,7 +195,7 @@ def __new_from_mlir_values__(self, values): class DSparkAttention(DSparkAttentionKernel): - """Attention over a 128-token rolling window and one 5/6-token draft block.""" + """Attention over a 128-token rolling window and one supported draft block.""" window_size = 128 block_size = 6 @@ -260,8 +260,9 @@ def __init__( ] if mismatches: raise ValueError("Unsupported DSpark kernel configuration: " + ", ".join(mismatches)) - if seq_len_q not in (5, 6): - raise ValueError(f"DSpark block size must be 5 or 6, got {seq_len_q}") + # Each query gets its own head tile; only draft storage bounds K. + if seq_len_q not in (1, 2, 3, 4, 5, 6, 7, 8): + raise ValueError(f"DSpark block size must be 1 through 8, got {seq_len_q}") if inverse_rope_dim not in (0, 64): raise ValueError(f"DSpark inverse_rope_dim must be 0 or 64, got {inverse_rope_dim}") diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index ccb4609d09c8..e05538bbd597 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -2103,13 +2103,23 @@ def write_context_windows_batched(self, main_hidden, positions, slots, mask, kv_ ) def _prepare_confidence_head_for_load(self, weights: Dict) -> None: - """Match the projection's optional bias to the checkpoint before loading.""" + """Prepare and validate the embedded confidence namespace before loading.""" head = self.dspark_model.mtp_layers[-1].confidence_head - bias_key = f"mtp_layers.{self.num_stages - 1}.confidence_head.proj.bias" - if head is not None and bias_key in weights and head.proj.bias is None: + prefix = f"mtp_layers.{self.num_stages - 1}.confidence_head." + module_weights = { + name[len(prefix) :]: value for name, value in weights.items() if name.startswith(prefix) + } + if head is None: + if module_weights: + raise ValueError("DSpark confidence head checkpoint has no constructed head") + return + if "proj.bias" in module_weights and head.proj.bias is None: # This runs before weight loading and graph capture. Empty storage # avoids consuming RNG or reinitializing the existing projection. head.proj.bias = nn.Parameter(head.proj.weight.new_empty(head.proj.out_features)) + # The generic V4 loader skips this parameter-free head container and + # visits its Linear child, so strict validation must run at this seam. + head._validate_checkpoint_weights(module_weights) def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): """Load the ``mtp.*`` draft weights from the (full) checkpoint dict. diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 569ea56dced9..52d3134b55d4 100644 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -579,6 +579,13 @@ def load_sts_temperatures(self, temperatures: torch.Tensor) -> None: def load_weights(self, weights: list) -> None: """Load every confidence-head checkpoint parameter strictly.""" (module_weights, ) = weights + self._validate_checkpoint_weights(module_weights) + for name, param in self.named_parameters(): + param.data.copy_(module_weights[name][:]) + + def _validate_checkpoint_weights( + self, module_weights: dict[str, torch.Tensor]) -> None: + """Validate a head namespace before generic embedded parameter loading.""" expected = dict(self.named_parameters()) unexpected = [key for key in module_weights if key not in expected] missing = [key for key in expected if key not in module_weights] @@ -593,7 +600,11 @@ def load_weights(self, weights: list) -> None: f"DSpark confidence head checkpoint is missing key(s) {sorted(missing)}" ) for name, param in expected.items(): - param.data.copy_(module_weights[name][:]) + actual_shape = tuple(module_weights[name].shape) + if actual_shape != tuple(param.shape): + raise ValueError( + f"DSpark confidence head checkpoint {name} has shape {actual_shape}; " + f"expected {tuple(param.shape)}") class Eagle3Attention(Attention): diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 2fde7afd6d5e..7fcdda9d4be5 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -85,12 +85,23 @@ from .seq_slot_manager import SeqSlotManager GB = 1 << 30 +_DEFAULT_INDEX_MAPPER_RESERVED_SLOTS = 1 +_CONFIDENCE_INDEX_MAPPER_RESERVED_SLOTS = 2 def ceil_div(a: int, b: int) -> int: return (a + b - 1) // b +def _get_num_reserved_index_slots(spec_config: Optional[SpeculativeConfig], + estimating_kv_cache: bool) -> int: + """Return IndexMapper headroom required by the execution policy.""" + if getattr(spec_config, "enable_confidence_scheduling", False): + return (0 if estimating_kv_cache else + _CONFIDENCE_INDEX_MAPPER_RESERVED_SLOTS) + return _DEFAULT_INDEX_MAPPER_RESERVED_SLOTS + + def _get_initial_lora_data_type( configured_lora_data_type: Optional[torch.dtype], ) -> Optional[torch.dtype]: @@ -3278,6 +3289,8 @@ def _create_kv_cache_manager( manager_extra_kwargs["max_copy_beam_width"] = max_copy_beam_width manager_extra_kwargs[ "cold_page_codec_provider"] = cold_page_codec_provider + manager_extra_kwargs["num_reserved_index_slots"] = ( + _get_num_reserved_index_slots(spec_config, estimating_kv_cache)) manager_extra_kwargs["kv_events_config"] = kv_events_config manager_extra_kwargs["joint_kv_cache_reuse"] = joint_kv_cache_reuse manager_extra_kwargs[ diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index f84f2c058026..6b55382bb149 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -349,6 +349,11 @@ def _is_mixed_encoder_decoder_batch(self, batch: ScheduledRequests) -> bool: and batch.num_generation_requests > 0) def _can_run_cuda_graph_batch(self, batch: ScheduledRequests) -> bool: + # A peer may be prefilling while this rank contains only generation. + # Host-authoritative windows use the common BREAKABLE body, never a + # rank-local full-generation graph or its device layout prologue. + if getattr(self, "_dspark_host_window_batch", None) is batch: + return False return (batch.can_run_cuda_graph or (self.enable_encoder_decoder_mixed_cuda_graph and self._is_mixed_encoder_decoder_batch(batch))) diff --git a/tensorrt_llm/_torch/pyexecutor/dspark_mixed_host_capacity.py b/tensorrt_llm/_torch/pyexecutor/dspark_mixed_host_capacity.py new file mode 100644 index 000000000000..ab49919ba7c1 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/dspark_mixed_host_capacity.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pure host mixed geometry with context/full-width capacity reservations. + +The caller must authenticate the current ready host survival matrix and retain +request ownership. This module neither reads device tensors nor prices routes. +Actual allocation is delegated to the existing DSpark top-k implementation. +""" + +from dataclasses import dataclass +from typing import Callable, Tuple + +import numpy as np +import torch + + +@dataclass(frozen=True) +class Capacity: + proposed_body_rows: int + eligible_rows: int + optional_budget: int + total_real_rows: int + expected_yield: float + + +@dataclass(frozen=True) +class LocalCurve: + request_ids: Tuple[int, ...] + survival: torch.Tensor + context_rows: int + fixed_generation_rows: int + physical_k: int + min_verify_len: int + native_expected_yield: float + candidates: Tuple[Capacity, ...] + + def candidate(self, body_rows: int) -> Capacity: + for item in self.candidates: + if item.proposed_body_rows == body_rows: + return item + raise ValueError("Unlisted or locally infeasible body capacity") + + +def _nonnegative_integer(value: int, name: str) -> None: + if type(value) is not int or value < 0: + raise ValueError(f"{name} must be a nonnegative integer") + + +def prepare_local_curve( + *, + survival: torch.Tensor, + request_ids: Tuple[int, ...], + context_rows: int, + fixed_generation_rows: int, + physical_k: int, + min_verify_len: int, + body_capacities: Tuple[int, ...], +) -> LocalCurve: + """Reserve context/full-width rows before filling eligible draft prefixes. + + Infeasible capacities are omitted. Context work contributes to paid rows, + not fictitious generation progress. Final-context progress is preserved by + the caller as a separate, candidate-independent component. + """ + for name, value in ( + ("context_rows", context_rows), + ("fixed_generation_rows", fixed_generation_rows), + ("physical_k", physical_k), + ("min_verify_len", min_verify_len), + ): + _nonnegative_integer(value, name) + if not 1 <= min_verify_len <= physical_k: + raise ValueError("Invalid draft prefix bounds") + if ( + not isinstance(survival, torch.Tensor) + or survival.device.type != "cpu" + or survival.ndim != 2 + or tuple(survival.shape) != (len(request_ids), physical_k) + ): + raise ValueError("Requires a ready CPU [eligible requests, physical K] matrix") + if any(type(item) is not int or item < 0 for item in request_ids) or len( + set(request_ids) + ) != len(request_ids): + raise ValueError("Request identities must be unique nonnegative integers") + if ( + not body_capacities + or any(type(item) is not int or item <= 0 for item in body_capacities) + or tuple(sorted(set(body_capacities))) != body_capacities + ): + raise ValueError("Body capacities must be positive, strictly increasing") + # Own the admitted host matrix: later copy-buffer reuse cannot change the + # capacity curve between the group vote and its allocation. + matrix = survival.detach().to(dtype=torch.float32).clone() + values = matrix.numpy().astype(np.float64, copy=False) + if ( + not np.isfinite(values).all() + or (values < 0).any() + or (values > 1).any() + or (np.diff(values, axis=1) > 0).any() + ): + raise ValueError("Survival must be finite probabilities and prefix-monotone") + num_real = len(request_ids) + floor_rows = num_real * (min_verify_len + 1) + native_rows = num_real * (physical_k + 1) + base_yield = float(num_real) + float(values[:, :min_verify_len].sum(dtype=np.float64)) + ordered = np.sort(values[:, min_verify_len:physical_k].reshape(-1))[::-1] + prefix = np.concatenate(([0.0], np.cumsum(ordered, dtype=np.float64))) + native_yield = base_yield + float(prefix[-1]) + reserve = context_rows + fixed_generation_rows + candidates = [] + for body_rows in body_capacities: + available = body_rows - reserve + if available < floor_rows: + continue + eligible_rows = min(available, native_rows) + optional_budget = eligible_rows - floor_rows + candidates.append( + Capacity( + body_rows, + eligible_rows, + optional_budget, + reserve + eligible_rows, + base_yield + float(prefix[optional_budget]), + ) + ) + return LocalCurve( + tuple(request_ids), + matrix, + context_rows, + fixed_generation_rows, + physical_k, + min_verify_len, + native_yield, + tuple(candidates), + ) + + +def allocate_candidate( + curve: LocalCurve, + body_rows: int, + original_allocate: Callable[[torch.Tensor, int, int, int], torch.Tensor], +): + """Use original DSpark top-k with survival_eps=0 and validate exact rows.""" + candidate = curve.candidate(body_rows) + lengths = original_allocate( + curve.survival, candidate.optional_budget, curve.physical_k, curve.min_verify_len + ) + if ( + not isinstance(lengths, torch.Tensor) + or lengths.device.type != "cpu" + or lengths.ndim != 1 + or len(lengths) != len(curve.request_ids) + ): + raise ValueError("Allocator returned incompatible request lengths") + if lengths.dtype not in (torch.int32, torch.int64): + raise ValueError("Allocator lengths must be integer") + windows = tuple(int(value) for value in lengths.tolist()) + if ( + any(not curve.min_verify_len <= value <= curve.physical_k for value in windows) + or sum(value + 1 for value in windows) != candidate.eligible_rows + ): + raise ValueError("Allocator did not spend the admitted eligible-row budget") + return tuple(zip(curve.request_ids, windows)) + + +def realized_paid_bucket( + curves: Tuple[LocalCurve, ...], proposed_body_rows: int, captured_body_rows: Tuple[int, ...] +) -> int: + """Canonicalize actual all-rank maximum; do not manufacture padding work.""" + if not curves or proposed_body_rows not in captured_body_rows: + raise ValueError("Requires a captured proposal and nonempty rank group") + if ( + any(type(item) is not int or item <= 0 for item in captured_body_rows) + or tuple(sorted(set(captured_body_rows))) != captured_body_rows + ): + raise ValueError("Captured body grid must be positive and strictly increasing") + needed = max(curve.candidate(proposed_body_rows).total_real_rows for curve in curves) + for bucket in captured_body_rows: + if bucket >= needed: + return bucket + raise ValueError("Actual rows exceed the captured body grid") diff --git a/tensorrt_llm/_torch/pyexecutor/dspark_mixed_host_snapshot.py b/tensorrt_llm/_torch/pyexecutor/dspark_mixed_host_snapshot.py new file mode 100644 index 000000000000..4a86a012f92c --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/dspark_mixed_host_snapshot.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Current ready host snapshot ownership; decode device mode is untouched.""" + +from dataclasses import dataclass +from typing import Callable, Tuple + +import torch + +from .dspark_mixed_host_capacity import LocalCurve, prepare_local_curve + + +@dataclass(frozen=True) +class AuthenticatedCurve: + curve: LocalCurve + request_slots: Tuple[Tuple[int, int], ...] + producing_iteration: int + snapshot_staged_iteration: int + staging_sequence: int + ranking_mode: str = "mixed_current_ready_host" + request_incarnations: Tuple[int, ...] = () + + +def prepare_current_mixed_curve( + *, + planner, + request_ids: Tuple[int, ...], + rows: Tuple[int, ...], + iteration: int, + staging_sequence: int, + context_rows: int, + fixed_generation_rows: int, + body_capacities: Tuple[int, ...], + original_compute_survival: Callable, + requests: Tuple[object, ...] = (), + incarnations: Tuple[int, ...] = (), +) -> AuthenticatedCurve | None: + """Return an owned current-host curve, or native fallback without waiting. + + This is an explicit mixed host-window contract, not fresh device ranking. + It uses the existing ready/event and request-addressed gather API even when + decode planner.device_windows is true. It never toggles that global flag. + Source/runtime admission and all-rank voting belong to the caller. + """ + if ( + type(iteration) is not int + or iteration < 0 + or type(staging_sequence) is not int + or staging_sequence < 0 + or len(request_ids) != len(rows) + ): + return None + k, floor = int(planner.max_verify_len), int(planner.cfg.min_verify_len) + if not request_ids: + # A context/dummy-only rank has no invented confidence ownership/yield. + curve = prepare_local_curve( + survival=torch.empty((0, k), dtype=torch.float32), + request_ids=(), + context_rows=context_rows, + fixed_generation_rows=fixed_generation_rows, + physical_k=k, + min_verify_len=floor, + body_capacities=body_capacities, + ) + return AuthenticatedCurve(curve, (), iteration, -1, staging_sequence) + meta = getattr(planner, "_mixed_current_snapshot_meta", None) + snapshot = planner._ready_snapshot() # Existing query(), never synchronize(). + if ( + snapshot is None + or meta is None + or meta.get("buffer") is not snapshot + or meta.get("event") is not planner._copy_event + or meta.get("staging_sequence") != staging_sequence + or type(meta.get("staged_iteration")) is not int + or not 0 <= meta["staged_iteration"] < iteration + or snapshot.device.type != "cpu" + ): + return None + owners = meta.get("owners") + if ( + not isinstance(owners, dict) + or len(set(request_ids)) != len(request_ids) + or len(set(rows)) != len(rows) + or any( + type(request_id) is not int + or request_id < 0 + or type(row) is not int + or row < 0 + or row >= snapshot.shape[0] + or owners.get(request_id) != row + for request_id, row in zip(request_ids, rows) + ) + ): + return None + # Ownership consistency above remains fail-closed. Freshness is a separate + # conservative partition: no writer proof means full K, never fake scores. + stamp = meta.get("confidence_stamp") + attempts = meta.get("producer_attempts") + expected_writer = staging_sequence - 1 # Copy precedes forward q. + valid = [] + if ( + expected_writer > 0 + and len(requests) == len(rows) + and len(incarnations) == len(rows) + and isinstance(attempts, dict) + and stamp is getattr(planner, "_host_confidence_stamp", None) + and isinstance(stamp, torch.Tensor) + and stamp.device.type == "cpu" + and stamp.dtype == torch.int32 + and stamp.ndim == 1 + and stamp.shape[0] == snapshot.shape[0] + ): + for index, (request_id, row, request, incarnation) in enumerate( + zip(request_ids, rows, requests, incarnations) + ): + attempt = attempts.get(request_id) + if ( + type(incarnation) is int + and incarnation > 0 + and type(attempt) is tuple + and len(attempt) == 3 + and attempt[0] is request + and type(attempt[1]) is int + and type(attempt[2]) is int + and attempt[1] == row + and attempt[2] == incarnation + and int(stamp[row]) == expected_writer + ): + valid.append(index) + fixed_generation_rows += (len(rows) - len(valid)) * (k + 1) + request_ids = tuple(request_ids[index] for index in valid) + rows = tuple(rows[index] for index in valid) + incarnations = tuple(incarnations[index] for index in valid) + if not rows: + curve = prepare_local_curve( + survival=torch.empty((0, k), dtype=torch.float32), + request_ids=(), + context_rows=context_rows, + fixed_generation_rows=fixed_generation_rows, + physical_k=k, + min_verify_len=floor, + body_capacities=body_capacities, + ) + return AuthenticatedCurve(curve, (), iteration, meta["staged_iteration"], staging_sequence) + selected = planner._gather_rows(num_gen_requests=len(rows), rows=rows, snapshot=snapshot) + if selected is None: + return None + survival = original_compute_survival(planner.apply_calibration(selected)) + try: + curve = prepare_local_curve( + survival=survival, + request_ids=tuple(request_ids), + context_rows=context_rows, + fixed_generation_rows=fixed_generation_rows, + physical_k=k, + min_verify_len=floor, + body_capacities=body_capacities, + ) + except ValueError: + return None + return AuthenticatedCurve( + curve, + tuple(zip(request_ids, rows)), + iteration, + meta["staged_iteration"], + staging_sequence, + request_incarnations=incarnations, + ) diff --git a/tensorrt_llm/_torch/pyexecutor/dspark_mixed_runtime.py b/tensorrt_llm/_torch/pyexecutor/dspark_mixed_runtime.py new file mode 100644 index 000000000000..3873b8adf2d7 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/dspark_mixed_runtime.py @@ -0,0 +1,474 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Admitted host-mixed allocation in the existing DSpark policy transaction. + +There is no prescribed action, diagnostic pricing mode, new collective, or +decode-table conversion. A missing quote/owner/route keeps mixed steps native. +""" + +import math +from dataclasses import dataclass, replace +from fractions import Fraction + +from ..speculative.dspark_mixed_evidence import AdmittedMixedCosts, MixedGeometry +from ..speculative.dspark_planner import _canonical_json_sha256 +from ..speculative.dspark_schedule import ( + DSparkScheduleConfig, + compute_survival, + schedule_verify_lens_topk, +) +from .dspark_mixed_host_capacity import allocate_candidate, prepare_local_curve +from .dspark_mixed_host_snapshot import AuthenticatedCurve, prepare_current_mixed_curve + +BODY_BUCKETS = (96, 128, 192, 256, 384, 512) +WIRE_WORDS = 46 +_MAGIC = 0x4D485031 +_VERSION = 1 +_K = 3 +_WORLD = 8 +_SCALE = 1_000_000 + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise ValueError(message) + + +def _integer(value: int) -> int: + _require(type(value) is int and 0 <= value < 2**63, "Invalid mixed geometry integer") + return value + + +def _membership(batch) -> tuple: + def rows(requests): + return tuple((id(r), r.py_request_id, r.py_seq_slot, bool(r.is_dummy)) for r in requests) + + return rows(batch.context_requests), rows(batch.generation_requests) + + +def _geometry(batch) -> tuple: + contexts = tuple( + ( + id(r), + r.context_chunk_size, + r.context_current_position, + r.py_num_compressed_tokens, + r.py_max_new_tokens, + bool(r.is_first_context_chunk), + ) + for r in batch.context_requests + ) + generations = tuple( + ( + id(r), + r.py_batch_idx, + len(r.py_draft_tokens), + r.get_num_tokens(0), + getattr(r.state, "name", None), + r.py_orig_prompt_len, + r.py_max_new_tokens, + ) + for r in batch.generation_requests + ) + return contexts, generations, tuple(id(r) for r in batch.context_requests_last_chunk) + + +def _remaining(request, runner) -> int: + return max( + 0, + min( + _integer(int(request.py_orig_prompt_len)) + _integer(int(request.py_max_new_tokens)), + runner._config.max_seq_len, + ) + - _integer(int(request.get_num_tokens(0))), + ) + + +def _greedy(batch) -> bool: + from tensorrt_llm.sampling_params import SamplingParams + + from .sampler.sampler_common import _request_get_sampling_params + + return all( + SamplingParams.params_imply_greedy_decoding( + temperature=p.temperature, + top_p=p.top_p, + top_k=p.top_k, + min_p=p.min_p, + top_p_decay=p.top_p_decay, + use_beam_search=p.use_beam_search, + ) + for p in (_request_get_sampling_params(r) for r in batch.all_requests()) + ) + + +@dataclass(frozen=True) +class _Ticket: + batch: object + membership: tuple + geometry: tuple + iteration: int + rank: int + eligible: tuple + curve: AuthenticatedCurve + vote: tuple[int, ...] + + +def stage_mixed_ownership(planner, worker, iteration: int, generation_requests) -> None: + """Bind the existing asynchronous copy, without reading device storage. + + An attempted forward is not a write. At the next decision, the copied + positive producer stamp must match that attempt and the current incarnation. + """ + if not planner._snapshot_valid or planner._host_buffer is None: + planner._mixed_current_snapshot_meta = None + return + previous = getattr(planner, "_mixed_current_snapshot_meta", None) + sequence = int(worker._draft_seq_host) + attempts = {} + for request in generation_requests: + if request.is_dummy: + continue + request_id = int(request.py_request_id) + slot = worker._req_to_slot.get(request_id) + incarnation = worker.confidence_incarnation_for(request_id) + if type(slot) is int and type(incarnation) is int and incarnation > 0: + attempts[request_id] = (request, slot, incarnation) + producer_attempts = ( + previous.get("generation_attempts", {}) + if isinstance(previous, dict) and previous.get("staging_sequence") == sequence - 1 + else {} + ) + planner._mixed_current_snapshot_meta = dict( + buffer=planner._host_buffer, + event=planner._copy_event, + owners=dict(worker._req_to_slot), + staged_iteration=int(iteration), + staging_sequence=sequence, + confidence_stamp=planner._host_confidence_stamp, + generation_attempts=attempts, + producer_attempts=producer_attempts, + ) + + +class MixedRuntime: + """A separately bound optional mixed path; decode admission is untouched.""" + + def __init__(self, runner, costs: AdmittedMixedCosts | None) -> None: + self.runner = runner + self.costs = costs + self._ticket: _Ticket | None = None + identity = ( + None + if costs is None + else { + "live": costs.identity.payload(), + "costs": costs.binding.costs_sha256, + "receipt": costs.binding.admission_sha256, + "protocol": costs.binding.protocol_sha256, + } + ) + digest = "0" * 64 if identity is None else _canonical_json_sha256(identity) + self._identity_words = tuple(int(digest[n : n + 8], 16) for n in range(0, 64, 8)) + + def reset_step(self) -> None: + self._ticket = None + self.runner._dspark_host_window_step = False + self.runner.cuda_graph_runner._dspark_host_window_batch = None + + def _route_ready(self, batch, planner) -> bool: + cfg, mapping = self.runner._config, self.runner.mapping + spec, breakable = cfg.spec_config, self.runner.breakable_cuda_graph_runner + return bool( + self.costs is not None + and self.runner._dspark_confidence_enabled + and self.runner._dspark_trims_submitted_tokens + and cfg.max_draft_len == spec.max_draft_len == spec.block_size == _K + and spec.decoding_type == "DSpark" + and spec.draft_is_embedded_in_target + and not getattr(spec, "use_rejection_sampling", False) + and mapping.tp_size == _WORLD + and mapping.pp_size == mapping.cp_size == 1 + and cfg.enable_attention_dp + and not cfg.torch_compile_enabled + and not self.runner.use_beam_search + and not cfg.is_multimodal + and not getattr(batch, "encoder_requests", ()) + and getattr(cfg.prefill_cuda_graph_backend, "name", None) == "BREAKABLE" + and tuple(sorted(cfg.prefill_cuda_graph_num_tokens)) == BODY_BUCKETS + and breakable is not None + and not breakable.is_capturing + and not breakable.is_warming_up + and all(breakable.has_graph(bucket) for bucket in BODY_BUCKETS) + and planner.max_verify_len == _K + and planner.cfg.min_verify_len == 1 + and _greedy(batch) + ) + + def local_vote( + self, batch, worker, planner, iteration: int, rank: int, *, executor_eligible: bool + ) -> list[int]: + """Append a fixed suffix to the existing allgather; never wait for a copy.""" + out = [ + _MAGIC, + _VERSION, + int(rank), + int(iteration), + sum(not r.is_dummy for r in batch.context_requests), + 0, + *self._identity_words, + *([0] * 32), + ] + _require(len(out) == WIRE_WORDS, "Mixed vote extent differs") + self._ticket = None + try: + _require( + executor_eligible and self._route_ready(batch, planner), "Unsupported mixed route" + ) + contexts, generations = list(batch.context_requests), list(batch.generation_requests) + # First-chunk KV preparation can resolve prefix-cache hits and + # change context position/chunk size after this policy transaction. + _require( + not any(r.is_first_context_chunk for r in contexts), + "Mixed context geometry is not finalized", + ) + native_rows = sum(_integer(r.context_chunk_size) for r in contexts) + native_rows += len(generations) * (_K + 1) + _require( + native_rows <= _integer(self.runner._config.max_num_tokens), + "Mixed native rows exceed the prepared token budget", + ) + real = [r for r in generations if not r.is_dummy] + _require( + all( + getattr(r.state, "name", None) not in (None, "GENERATION_COMPLETE") + and _remaining(r, self.runner) > 0 + for r in real + ), + "Completed generation", + ) + meta = getattr(planner, "_mixed_current_snapshot_meta", None) + owners = meta.get("owners") if isinstance(meta, dict) else None + prospective = [ + r for r in real if r.py_batch_idx is not None and len(r.py_draft_tokens) >= _K + ] + _require( + len({r.py_request_id for r in prospective}) == len(prospective), "Duplicate owner" + ) + eligible = [ + r for r in prospective if not isinstance(owners, dict) or r.py_request_id in owners + ] + fixed = [r for r in generations if id(r) not in {id(e) for e in eligible}] + chunks = tuple(_integer(r.context_chunk_size) for r in contexts) + _require(all(chunks), "Empty context chunk") + curve = prepare_current_mixed_curve( + planner=planner, + request_ids=tuple(int(r.py_request_id) for r in eligible), + rows=tuple(worker.confidence_row_for(r.py_request_id) for r in eligible), + iteration=iteration, + staging_sequence=int(worker._draft_seq_host), + context_rows=sum(chunks), + fixed_generation_rows=len(fixed) * (_K + 1), + body_capacities=BODY_BUCKETS, + original_compute_survival=compute_survival, + requests=tuple(eligible), + incarnations=tuple( + worker.confidence_incarnation_for(r.py_request_id) for r in eligible + ), + ) + _require(curve is not None, "Current owned snapshot unavailable") + eligible = [r for r in eligible if r.py_request_id in curve.curve.request_ids] + fixed = [r for r in generations if id(r) not in {id(e) for e in eligible}] + useful = curve.curve.survival.clone() + for row, request in enumerate(eligible): + useful[row, max(0, _remaining(request, self.runner) - 1) :] = 0 + curve = replace( + curve, + curve=prepare_local_curve( + survival=useful, + request_ids=curve.curve.request_ids, + context_rows=sum(chunks), + fixed_generation_rows=len(fixed) * (_K + 1), + physical_k=_K, + min_verify_len=1, + body_capacities=BODY_BUCKETS, + ), + ) + real_fixed = [r for r in fixed if not r.is_dummy] + last = {id(r) for r in batch.context_requests_last_chunk if not r.is_dummy} + _require( + last <= {id(r) for r in contexts if not r.is_dummy}, + "Final-context ownership differs", + ) + context_progress = sum(r.py_max_new_tokens > 0 for r in contexts if id(r) in last) + lower = context_progress + len(real_fixed) + upper = context_progress + sum( + min(_K + 1, _remaining(r, self.runner)) for r in real_fixed + ) + histories = [_integer(int(r.get_num_tokens(0))) for r in real] + out[14:26] = [ + sum(chunks), + sum(n * n for n in chunks), + sum(n * _integer(r.context_current_position) for n, r in zip(chunks, contexts)), + sum(n * _integer(r.py_num_compressed_tokens) for n, r in zip(chunks, contexts)), + sum(histories), + max(histories, default=0), + len(real_fixed), + out[4], + len(eligible), + len(contexts), + len(generations), + len(fixed) * (_K + 1), + ] + out[26:28] = [ + math.ceil((curve.curve.native_expected_yield + upper) * _SCALE), + native_rows, + ] + for index, bucket in enumerate(BODY_BUCKETS): + try: + cell = curve.curve.candidate(bucket) + except ValueError: + continue + out[28 + 3 * index : 31 + 3 * index] = [ + 1, + math.floor((cell.expected_yield + lower) * _SCALE), + cell.total_real_rows, + ] + out[5] = 1 + _require(all(type(n) is int and 0 <= n < 2**63 for n in out), "Vote integer overflow") + self._ticket = _Ticket( + batch, + _membership(batch), + _geometry(batch), + iteration, + rank, + tuple(eligible), + curve, + tuple(out), + ) + except ( + ValueError, + TypeError, + KeyError, + AttributeError, + RuntimeError, + IndexError, + OverflowError, + ): + out[5] = 0 + self._ticket = None + # Invalid optional geometry cannot poison the collective's int64 + # extent. Its real-context count still selects group-native mixed. + out[14:] = [0] * (WIRE_WORDS - 14) + return out + + def split_votes(self, payloads, legacy_words: int) -> tuple[list, list]: + _require( + len(payloads) == _WORLD and all(len(p) == legacy_words + WIRE_WORDS for p in payloads), + "Mixed policy suffix shape differs", + ) + peers = [list(p[legacy_words:]) for p in payloads] + for rank, row in enumerate(peers): + _require( + all(type(n) is int and 0 <= n < 2**63 for n in row) + and row[:3] == [_MAGIC, _VERSION, rank] + and row[5] in (0, 1), + "Mixed policy header differs", + ) + for index, bucket in enumerate(BODY_BUCKETS): + feasible, amount, total = row[28 + 3 * index : 31 + 3 * index] + _require( + feasible in (0, 1) + and (feasible or amount == total == 0) + and (not feasible or total <= bucket and amount <= row[26]), + "Mixed capacity vote differs", + ) + _require(len({r[3] for r in peers}) == 1, "Mixed iteration differs") + return [list(p[:legacy_words]) for p in payloads], peers + + def apply(self, batch, worker, planner, peers, iteration: int) -> bool: + """Return False only for all-decode; preserve its original exact policy.""" + if not sum(row[4] for row in peers): + return False + _require(all(row[3] == iteration for row in peers), "Mixed publication iteration differs") + chosen, windows = None, {} + if self.costs is not None and all( + row[5] and tuple(row[6:14]) == self._identity_words for row in peers + ): + features = tuple(tuple(row[14:26]) for row in peers) + native_rows = tuple(row[27] for row in peers) + native_yield = sum(row[26] for row in peers) + for index, capacity in enumerate(BODY_BUCKETS): + cells = [row[28 + 3 * index : 31 + 3 * index] for row in peers] + if not all(cell[0] for cell in cells): + continue + try: + geometry = MixedGeometry( + features, capacity, tuple(c[2] for c in cells), native_rows + ) + except ValueError: + continue + quote = self.costs.lookup(geometry) + if quote is None: + continue + candidate = Fraction(sum(cell[1] for cell in cells), 1) / Fraction( + str(quote.candidate_upper_ms) + ) + native = Fraction(native_yield, 1) / Fraction(str(quote.native_lower_ms)) + margin = 1 + Fraction(str(self.costs.minimum_predicted_gain)) + if candidate > native and candidate >= native * margin: + item = (candidate, -capacity, geometry) + if chosen is None or item[:2] > chosen[:2]: + chosen = item + if chosen is not None: + ticket = self._ticket + _require( + ticket is not None + and ticket.batch is batch + and ticket.iteration == iteration + and ticket.membership == _membership(batch) + and ticket.geometry == _geometry(batch) + and ticket.vote == tuple(peers[ticket.rank]), + "Mixed request/geometry ticket changed", + ) + _require( + worker._draft_seq_host == ticket.curve.staging_sequence + and tuple(worker.confidence_row_for(r.py_request_id) for r in ticket.eligible) + == tuple(row for _, row in ticket.curve.request_slots) + and tuple( + worker.confidence_incarnation_for(r.py_request_id) for r in ticket.eligible + ) + == ticket.curve.request_incarnations, + "Mixed snapshot/slot incarnation changed", + ) + + def original_allocate(survival, budget, k, floor): + return schedule_verify_lens_topk( + survival=survival, + budget=budget, + cfg=DSparkScheduleConfig(k, floor, survival_eps=0.0), + ) + + windows = dict( + allocate_candidate(ticket.curve.curve, chosen[2].capacity, original_allocate) + ) + actual = ticket.curve.curve.context_rows + ticket.curve.curve.fixed_generation_rows + actual += sum(value + 1 for value in windows.values()) + _require( + actual == chosen[2].actual_rank_rows[ticket.rank], "Mixed allocated rows differ" + ) + values = [ + _K if r.is_dummy else windows.get(r.py_request_id, _K) + for r in batch.generation_requests + ] + _require(all(type(v) is int and 1 <= v <= _K for v in values), "Mixed window bounds differ") + for request, value in zip(batch.generation_requests, values): + request.py_verify_len = value + graph = self.runner.cuda_graph_runner + graph.agreed_ragged_bucket = None + graph.ragged_pad_verify_len = 0 + graph.ragged_zero_real_high_rows = 0 + graph._dspark_host_window_batch = batch + self.runner._dspark_device_budget = None + self.runner._dspark_host_window_step = True + return True diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 49ee7677e181..2d25b76c30d1 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -530,26 +530,68 @@ def _init_decoder_state(self) -> None: self._dspark_confidence_enabled = bool( spec_config is not None and getattr(spec_config, "enable_confidence_scheduling", False) ) + self._dspark_fused_scheduler_enabled = bool( + spec_config is not None + and getattr(spec_config, "enable_fused_confidence_scheduler", 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_device_window_workspace = None + self._dspark_host_window_step = False + self._dspark_mixed_runtime = None + mixed_config = getattr(spec_config, "confidence_mixed_evidence", None) + if ( + self._dspark_confidence_enabled + and mixed_config is not None + and self._config.max_draft_len == 3 + and self.mapping.tp_size == 8 + and self.mapping.pp_size == self.mapping.cp_size == 1 + and self._config.enable_attention_dp + ): + from tensorrt_llm._torch.pyexecutor.dspark_mixed_runtime import MixedRuntime + from tensorrt_llm._torch.speculative.dspark_mixed_evidence import load_mixed_evidence + + resolution = load_mixed_evidence(mixed_config) + self._dspark_mixed_runtime = MixedRuntime(self, resolution.costs) + logger.info(f"DSpark mixed evidence: {resolution.reason}") + if self._dspark_device_windows: + from tensorrt_llm._torch.speculative.dspark_select import DeviceWindowWorkspace + + self._dspark_device_window_workspace = DeviceWindowWorkspace.allocate( + max_rows=self._config.max_batch_size + 1, + max_tokens=self._config.max_num_tokens, + device=self.input_ids_cuda.device, + ) self._dspark_sps_cost_table = None + self._dspark_sps_artifact_metadata = None + self._dspark_sps_observed_runtime = 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_sps_cost_table, self._dspark_sps_artifact_metadata = ( + 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 + ), + validate_freshness=spec_config.validate_confidence_artifact_freshness, + ) ) + self._dspark_sps_observed_runtime = { + "max_draft_len": int(spec_config.max_draft_len), + "rank_local_graph_batch_sizes": list(self._config.cuda_graph_batch_sizes), + "gpu": torch.cuda.get_device_name(self.input_ids_cuda.device), + "torch_version": torch.__version__, + "validate_artifact_freshness": spec_config.validate_confidence_artifact_freshness, + } 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 @@ -928,6 +970,7 @@ def _warmup_scheduled(self, resource_manager: ResourceManager, kv_cache_manager) with self._warmup_timer.phase("dsa_prewarm"): self._ensure_dsa_attn_metadata_for_warmup(resource_manager) self._warmup_dspark_ragged_compressor_metadata() + self._warmup_dspark_fused_scheduler() 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") @@ -2343,7 +2386,8 @@ def _capture_variant( # Dummy warmup requests carry no sampling params, so the greedy pass # needs no override while the advanced one has to force the flag. _capture_variant("greedy") - _capture_variant("advanced sampling", force_non_greedy=True) + if self._supports_advanced_draft_sampling(): + _capture_variant("advanced sampling", force_non_greedy=True) else: # TorchSampler branch: FULL carries no sampling in the graph and is # what every batch replays unless in-graph sampling is enabled, so it @@ -3136,7 +3180,9 @@ def _sync_group_all_greedy_sample(self, spec_metadata) -> None: Why the sampling-path choice must be group-uniform under ADP + LM-head TP is documented on the anchor, - ``SpecMetadata.group_all_greedy_sample``. Local contract: called once + ``SpecMetadata.group_all_greedy_sample``. Confidence also needs a + uniform flag to reject unsupported advanced sampling on every rank. + Local contract: called once per iteration, right after ``update_is_all_greedy_sample`` and BEFORE the CUDA graph key is built. The gate is pure config (identical on every rank), so ranks also agree on whether the exchange happens; the @@ -3145,19 +3191,40 @@ def _sync_group_all_greedy_sample(self, spec_metadata) -> None: ``all_rank_num_tokens`` exchange, which runs in ``_prepare_inputs`` -- after the graph key, too late for the key to see the synced value. """ - # enable_lm_head_tp_in_adp implies enable_attention_dp (asserted in - # Mapping.__init__), so ADP needs no separate check here. - if not (self.mapping.enable_lm_head_tp_in_adp and spec_metadata.use_rejection_sampling): + # Confidence's unsupported advanced path must also be rejected by all + # ranks before graph replay, even without LM-head TP in attention DP. + # Feature-off engines retain the existing LM-head-TP gate. + if not ( + spec_metadata.use_rejection_sampling + and (self.mapping.enable_lm_head_tp_in_adp or self._dspark_confidence_enabled) + ): return local_flag = bool(spec_metadata.is_all_greedy_sample) - all_flags = self.dist.tp_allgather_int64([local_flag])[:, 0] - spec_metadata.group_all_greedy_sample = bool(all_flags.all()) + if self.mapping.tp_size > 1: + all_flags = self.dist.tp_allgather_int64([local_flag])[:, 0] + spec_metadata.group_all_greedy_sample = bool(all_flags.all()) + else: + spec_metadata.group_all_greedy_sample = local_flag # Also overwrite the live flag directly: this iteration's scan already # ran (update_is_all_greedy_sample just returned) and the CUDA graph # key reads the flag next -- the stored override only takes effect on # the NEXT rescan (populate), which is after key selection. spec_metadata.is_all_greedy_sample = spec_metadata.group_all_greedy_sample + def _supports_advanced_draft_sampling(self) -> bool: + """Whether the engine may capture an advanced draft-sampling variant.""" + return not ( + self._dspark_confidence_enabled and self._config.spec_config.use_rejection_sampling + ) + + def _validate_dspark_draft_sampling(self, spec_metadata: SpecMetadata) -> None: + """Reject an unsupported group-uniform sampler path before graph lookup.""" + if self._dspark_confidence_enabled and spec_metadata.wants_advanced_draft_sampling: + raise ValueError( + "DSpark confidence scheduling does not support advanced rejection draft sampling: " + "confidence is conditioned on the greedy proposal prefix" + ) + def _is_final_multimodal_context_decode_compatible(self, request: LlmRequest) -> bool: """Return whether the final prompt token uses the decode input path. @@ -3218,6 +3285,90 @@ def _record_cached_kv_tokens_per_req( num_cached_tokens_per_seq ) - sum(counts) + @torch.inference_mode() + def _warmup_dspark_fused_scheduler(self) -> None: + """Compile the optional policy-neutral scheduler before live requests. + + The device-window prologue is deliberately skipped during whole-model + CUDA-graph capture, so its Triton rank/finalize kernels would otherwise + compile on the first live compact step. Compile every authenticated + exact-SPS graph row count on every rank and record only successful + shapes. Unsupported shapes retain the established tensor fallback. + """ + self._dspark_fused_schedule_ready_sizes = set() + if ( + not self._dspark_fused_scheduler_enabled + or not self._dspark_confidence_enabled + or not self._dspark_trims_submitted_tokens + or not self._dspark_device_windows + ): + return + + worker = self._get_spec_worker() + planner = getattr(worker, "verify_planner", None) + exact_table = planner.exact_cost_table if planner is not None else None + if exact_table is None: + return + + from tensorrt_llm._torch.speculative.dspark_schedule import ( + schedule_verify_lens_topk_fused_fill, + ) + + cfg = planner.cfg + if cfg.survival_eps <= 0.0: + logger.warning_once( + "DSpark fused scheduler requires survival_eps > 0; using the " + "tensor scheduler for every graph size", + key="dspark_fused_scheduler_zero_epsilon", + ) + return + ready = [] + with torch.inference_mode(): + for graph_bs in sorted(int(g) for g in exact_table.tables): + if graph_bs > 256: + logger.warning_once( + "DSpark fused scheduler supports at most " + f"256 rows; G={graph_bs} will use the tensor fallback", + key=(f"dspark_fused_scheduler_unsupported_{graph_bs}"), + ) + continue + survival = torch.ones( + (graph_bs, int(cfg.block_size)), + dtype=torch.float32, + device="cuda", + ) + try: + warmup_budget = min(1, graph_bs * int(cfg.schedulable_per_request)) + graph_num_tokens = graph_bs * (int(cfg.min_verify_len) + 1) + warmup_budget + schedule_verify_lens_topk_fused_fill( + survival=survival, + budget=warmup_budget, + num_real=graph_bs, + pad_len=1, + cfg=cfg, + graph_num_tokens=graph_num_tokens, + ) + except Exception as exc: + logger.warning_once( + "DSpark fused scheduler prewarm failed for " + f"G={graph_bs}; using the tensor fallback. " + f"{type(exc).__name__}: {exc}", + key=(f"dspark_fused_scheduler_prewarm_failure_{graph_bs}"), + ) + continue + try: + torch.cuda.synchronize() + except Exception as exc: + # An asynchronous kernel failure may poison the context; + # continuing with the tensor fallback is not safe. + raise RuntimeError( + f"DSpark fused scheduler failed its startup CUDA gate for G={graph_bs}" + ) from exc + ready.append(graph_bs) + + self._dspark_fused_schedule_ready_sizes = set(ready) + logger.info(f"DSpark fused scheduler prewarm complete for graph batch sizes {ready}") + @torch.inference_mode() def _warmup_dspark_ragged_compressor_metadata(self) -> None: """Compile DSv4 ragged metadata for a non-CUDA-graph batch size. @@ -3580,6 +3731,41 @@ def fit_ragged_verify_lens( request.py_verify_len = int(tokens) - 1 return int(bucket) + def _select_dspark_windows_with_fused_fallback( + self, + *, + select_fn: Callable[..., Any], + selector_kwargs: dict[str, Any], + padded_bs: int, + ) -> Any: + """Run the optional scheduler once, retiring only a failed graph G.""" + from tensorrt_llm._torch.speculative.dspark_schedule import DSparkFusedScheduleError + + ready_sizes = getattr(self, "_dspark_fused_schedule_ready_sizes", set()) + use_fused_exact = padded_bs in ready_sizes + try: + return select_fn(**selector_kwargs, use_fused_exact=use_fused_exact) + except DSparkFusedScheduleError as exc: + if not use_fused_exact: + raise + self._retire_dspark_fused_shape(padded_bs, exc, "scheduler") + return select_fn(**selector_kwargs, use_fused_exact=False) + + def _retire_dspark_fused_shape(self, padded_bs: int, exc: Exception, component: str) -> None: + """Permanently route one failed graph-row shape to the tensor oracle.""" + ready_sizes = set(getattr(self, "_dspark_fused_schedule_ready_sizes", set())) + ready_sizes.discard(padded_bs) + self._dspark_fused_schedule_ready_sizes = ready_sizes + failure_counts = getattr(self, "_dspark_fused_schedule_failure_counts", {}) + failure_counts[padded_bs] = failure_counts.get(padded_bs, 0) + 1 + self._dspark_fused_schedule_failure_counts = failure_counts + logger.warning_once( + f"DSpark fused {component} failed for G={padded_bs}; disabling " + "that shape and rerunning the exact tensor fallback. " + f"{type(exc).__name__}: {exc}", + key=f"dspark_fused_{component}_runtime_failure_{padded_bs}", + ) + def _apply_device_window_prologue(self, inputs, new_tensors_device) -> bool: """Re-rank this step's verify windows on device, with fresh confidence. @@ -3599,12 +3785,16 @@ def _apply_device_window_prologue(self, inputs, new_tensors_device) -> bool: """ from tensorrt_llm._torch.speculative.dspark_select import ( gather_packed_draft_tokens, + materialize_compact_layout, select_windows_device, ) + from tensorrt_llm._torch.speculative.dspark_schedule import DSparkFusedScheduleError runner = self.cuda_graph_runner budget = self._dspark_device_budget self._dspark_device_budget = None + if self._dspark_host_window_step: + return False if budget is None or runner.agreed_ragged_bucket is None: return False if not getattr(self, "_dspark_prev_covers_batch", False): @@ -3649,18 +3839,12 @@ def _apply_device_window_prologue(self, inputs, new_tensors_device) -> bool: # 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.ragged_verify_lens_cuda qo_buf = self.ragged_qo_indptr_cuda - split_lens = lens_buf[:padded_bs].clone() - split_qo = qo_buf[: padded_bs + 1].to(torch.long) - past_seen = self.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( + selector_kwargs = dict( confidence_logits=worker.staged_confidence_buffer(), slot_idx=worker.batch_slot_view(padded_bs), num_real=n_real, @@ -3671,26 +3855,67 @@ def _apply_device_window_prologue(self, inputs, new_tensors_device) -> bool: stamp=stamps, expected_stamp=expected_stamp, pad_len=pad_len_tok, + workspace=getattr(self, "_dspark_device_window_workspace", None), + ) + result = self._select_dspark_windows_with_fused_fallback( + select_fn=select_windows_device, + selector_kwargs=selector_kwargs, + padded_bs=padded_bs, ) + fused_snapshot_valid = False + if result.workspace is not None: + try: + materialize_compact_layout( + result=result, + old_verify_lens=lens_buf[:padded_bs], + old_qo_indptr=qo_buf[: padded_bs + 1], + batch_slots=self.previous_batch_indices_cuda[:n_real], + new_tokens=new_tensors_device.new_tokens, + new_tokens_lens=new_tensors_device.new_tokens_lens, + next_draft_tokens=new_tensors_device.next_draft_tokens, + input_ids=self.input_ids_cuda[:bucket], + position_ids=self.position_ids_cuda[:bucket], + previous_pos_indices=self.previous_pos_indices_cuda[:real_tokens], + previous_pos_id_offsets=self.previous_pos_id_offsets_cuda[:real_tokens], + draft_tokens=self.draft_tokens_cuda[: real_tokens - n_real], + kv_lens=attn_metadata.kv_lens_cuda[:padded_bs], + previous_kv_lens_offsets=self.previous_kv_lens_offsets_cuda[:padded_bs], + num_real=n_real, + real_tokens=real_tokens, + ) + except DSparkFusedScheduleError as exc: + # Row state commits in the final kernel, so a synchronous + # launch failure before it is safe to rerun from the original + # shape split. Every token output is overwritten below. + self._retire_dspark_fused_shape(padded_bs, exc, "layout materializer") + fused_snapshot_valid = bool(getattr(exc, "past_seen_valid", False)) + result = select_windows_device(**selector_kwargs, use_fused_exact=False) + else: + req_idx = result.req_idx + spec_metadata.remap_expanded_sampling_params(req_idx, bucket) + apply_device_layout(result.verify_lens, req_idx, result.kv_correction) + return True + + # Tensor oracle/fallback: keep the previous independently implemented + # materialization unchanged. It also repairs every token output if an + # optional fused launch failed synchronously above. + split_lens = lens_buf[:padded_bs].clone() + split_qo = qo_buf[: padded_bs + 1].to(torch.long) + if fused_snapshot_valid: + # The first fused launch saved these values before any token write. + past_seen = self._dspark_device_window_workspace.past_seen[:padded_bs] + else: + past_seen = self.position_ids_cuda[split_qo[:-1]].clone() + 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. + # Host staging includes the window in both num_cached and seq_lens_kv: + # kv_lens = past + 2*shape_lens. Fresh windows therefore require twice + # the row delta here and its negative in the replay overlap offset. 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. + # After replay adds its offset, the sum is past + fresh_window + new_len. self.previous_kv_lens_offsets_cuda[:padded_bs] -= window_delta.to( self.previous_kv_lens_offsets_cuda.dtype ) @@ -3722,10 +3947,8 @@ def _apply_device_window_prologue(self, inputs, new_tensors_device) -> bool: self.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. + # Draft-token storage excludes every bonus/anchor position, including + # the last row's anchor when the allocation is full. real_draft = real_tokens - n_real if real_draft > 0: self.draft_tokens_cuda[:real_draft].copy_( @@ -3756,7 +3979,12 @@ def _ragged_token_lens(generation_requests) -> Optional[List[int]]: return [1 + int(v) for v in verify_lens] def _publish_gen_token_layout( - self, attn_metadata, generation_requests, *, runtime_draft_len: int + self, + attn_metadata, + generation_requests, + *, + runtime_draft_len: int, + host_authoritative: bool = False, ) -> None: """Hand the attention metadata this step's gen-token layout, before it prepares. @@ -3787,7 +4015,9 @@ def _publish_gen_token_layout( # 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 + self._dspark_trims_submitted_tokens + and self._dspark_device_windows + and not host_authoritative ) def _pinned_host(self, key: str, values, dtype) -> torch.Tensor: @@ -5065,6 +5295,7 @@ def previous_seq_slots_device(): attn_metadata, scheduled_requests.generation_requests, runtime_draft_len=runtime_draft_len, + host_authoritative=bool(getattr(self, "_dspark_host_window_step", False)), ) refresh_seq_lens = not attn_metadata.is_cuda_graph if ( @@ -5618,6 +5849,7 @@ def _forward_decoder( if spec_metadata is not None: spec_metadata.update_is_all_greedy_sample(padded_graph_requests.all_requests()) self._sync_group_all_greedy_sample(spec_metadata) + self._validate_dspark_draft_sampling(spec_metadata) peft_cache_data_type = None if self._lora.cuda_graph_manager is not None: diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index f064391d0037..c14b982fefdd 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -840,6 +840,13 @@ def _dspark_device_budget(self): return None return getattr(self._runner, "_dspark_device_budget", None) + @property + def _dspark_mixed_runtime(self): + """Optional mixed controller owned by the actual decoder runner.""" + if not isinstance(self._runner, DecoderRunner): + return None + return self._runner._dspark_mixed_runtime + @_dspark_device_budget.setter def _dspark_device_budget(self, value) -> None: if isinstance(self._runner, DecoderRunner): diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index c51ec2bc75e7..033fcd972890 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -79,6 +79,7 @@ from .adp_iter_stats import ADPIterStatsBuffer from .connectors.kv_cache_connector import KvCacheConnectorManager from .connectors.kv_cache_layout import build_kv_cache_layout_v2 +from .cuda_graph_runner import ADPShapeAgreement from .disagg_adapter import PyExecutorEffects, PyExecutorRequestRegistry from .dwdp import DwdpManager from .error_classification import ErrorBudget @@ -133,6 +134,150 @@ _UNBOUNDED_STATS_MAX_LEN = -1 _V2_RETENTION_IGNORED_LOG_KEY = "kv_cache_manager_v2_retention_config_ignored" +_DSPARK_EXACT_WIRE_PREFIX_LEN = 16 +_DSPARK_ADP_WIRE_TRAILER_LEN = 3 +_DSPARK_EXACT_NATIVE_YIELD_INDEX = _DSPARK_EXACT_WIRE_PREFIX_LEN - 1 + + +def _encode_dspark_exact_expected_yields( + *, exact_local, num_cells: int, + yield_scale: int) -> Tuple[int, Tuple[int, ...]]: + """Encode one rank's exact yields without splitting the collective.""" + if yield_scale <= 0: + raise ValueError("DSpark exact-policy yield scale must be positive") + num_cells = int(num_cells) + if num_cells < 0: + raise ValueError("DSpark exact-policy cell count must be nonnegative") + if exact_local is None: + return 0, tuple(-1 for _ in range(num_cells)) + + native = float(exact_local.native_expected_yield) + native_valid = math.isfinite(native) and native >= 0.0 + native_wire = int(math.ceil(native * yield_scale)) if native_valid else 0 + local_values = tuple(exact_local.compact_expected_yields) + compact_wires = [] + for cell_index in range(num_cells): + value = (float(local_values[cell_index]) + if cell_index < len(local_values) else math.nan) + if (not native_valid or not math.isfinite(value) or value < 0.0 + or value > native): + compact_wires.append(-1) + else: + compact_wires.append(int(math.floor(value * yield_scale))) + return native_wire, tuple(compact_wires) + + +def _decode_dspark_exact_expected_yields( + *, + payloads: List[List[int]], + exact_cells: Iterable[Tuple[int, int]], + graph_batch_size: int, + yield_scale: int, +) -> Tuple[float, Dict[int, float], Dict[int, float]]: + """Decode aggregate yields and worst-rank loss per real request.""" + if yield_scale <= 0: + raise ValueError("DSpark exact-policy yield scale must be positive") + native_expected_yield = sum( + float(payload[_DSPARK_EXACT_NATIVE_YIELD_INDEX]) + for payload in payloads) / yield_scale + compact_expected_yields = {} + compact_max_yield_losses_per_request = {} + for cell_index, (cell_graph_batch_size, + verifier_budget) in enumerate(exact_cells): + if cell_graph_batch_size != graph_batch_size: + continue + rank_yields = [ + float(payload[_DSPARK_EXACT_WIRE_PREFIX_LEN + cell_index]) + for payload in payloads + ] + infeasible = any(value < 0.0 for value in rank_yields) + compact_expected_yields[verifier_budget] = (0.0 if infeasible else + sum(rank_yields) / + yield_scale) + active_rank_losses = [] + if not infeasible: + for payload, compact_yield_wire in zip(payloads, rank_yields): + num_real_requests = int(payload[1]) + if num_real_requests <= 0: + continue + native_yield_wire = float( + payload[_DSPARK_EXACT_NATIVE_YIELD_INDEX]) + if compact_yield_wire > native_yield_wire: + raise RuntimeError( + "DSpark exact-policy compact yield exceeds native yield" + ) + active_rank_losses.append( + (native_yield_wire - compact_yield_wire) / yield_scale / + num_real_requests) + compact_max_yield_losses_per_request[verifier_budget] = max( + active_rank_losses, default=0.0) + return ( + native_expected_yield, + compact_expected_yields, + compact_max_yield_losses_per_request, + ) + + +def _validate_dspark_exact_bucket(*, exact_shape, bucket) -> None: + """Require the graph key to match the globally selected exact cell.""" + if exact_shape is None: + return + expected_bucket = int(exact_shape[1]) + if bucket is None or int(bucket) != expected_bucket: + raise RuntimeError( + "DSpark exact policy did not publish its selected captured " + f"verifier bucket V={expected_bucket}; got {bucket}") + + +def _validate_dspark_adp_acceptance_gate( + *, confidence_enabled: bool, attention_dp_enabled: bool, + speculation_gate_enabled: bool) -> None: + """Reject a rank-local gate that can split the ADP collective protocol.""" + if confidence_enabled and attention_dp_enabled and speculation_gate_enabled: + raise ValueError( + "DSpark confidence scheduling with attention DP does not support " + "acceptance_rate_window_size/acceptance_rate_threshold because " + "the per-rank gate can break collective ordering") + + +def _classify_dspark_exact_generation_rows( + generation_requests: List[LlmRequest], +) -> Tuple[Optional[List[LlmRequest]], bool]: + """Return logical exact rows and whether one idle ADP dummy is scheduled.""" + if not generation_requests: + return None, False + real_requests = [ + request for request in generation_requests if not request.is_dummy + ] + if len(real_requests) == len(generation_requests): + return real_requests, False + one_adp_dummy = bool(not real_requests and len(generation_requests) == 1 + and generation_requests[0].is_attention_dp_dummy) + if one_adp_dummy: + return real_requests, True + return None, False + + +def _dspark_exact_common_graph_batch_size( + payloads: List[List[int]], round_up: Callable[[int], int]) -> int: + """Resolve a common exact G, preserving zero as all-rank termination.""" + widest_rows = max((int(payload[1]) for payload in payloads), default=0) + return 0 if widest_rows <= 0 else int(round_up(widest_rows)) + + +def _dspark_exact_secondary_padding_ready( + payloads: List[List[int]], + *, + graph_batch_size: int, + verifier_budget: int, + secondary_ready_index: int, +) -> bool: + """Whether every zero-real peer can represent the selected remainder.""" + if verifier_budget % graph_batch_size <= 1: + return True + return all( + int(payload[1]) != 0 or bool(int(payload[secondary_ready_index])) + for payload in payloads) class _ADPForwardIntent(IntEnum): @@ -639,6 +784,12 @@ def __init__( threshold = getattr(spec_config, 'acceptance_rate_threshold', None) if window and threshold is not None: self.speculation_gate = SpeculationGate(window, threshold) + _validate_dspark_adp_acceptance_gate( + confidence_enabled=getattr(self.model_engine, + "_dspark_confidence_enabled", False), + attention_dp_enabled=self.enable_attention_dp, + speculation_gate_enabled=self.speculation_gate is not None, + ) # response used data self.response_lock = threading.Lock() @@ -2090,7 +2241,7 @@ def get_latest_iteration_stats(self): Returns the per-iterations statistics computed since last call to this method. Contains at most iter_stats_max_iterations iterations. """ - if self.enable_iter_perf_stats == False: + if not self.enable_iter_perf_stats: return [] latest_stats = (IterationStats(), None) @@ -3931,13 +4082,414 @@ def _drain_relay_sends_before_forward(self, microbatch_id: int) -> None: self.dist.pp_size) % self.num_micro_batches self.wait_on_pp_send_handles(self.send_handles, due_microbatch_id) + def _stage_dspark_confidence(self, worker, planner, gen_requests, + mixed) -> None: + buffer = worker.staged_confidence_buffer() + if buffer is None: + return + worker.bump_draft_seq() + if mixed is not None and mixed.costs is not None: + from .dspark_mixed_runtime import stage_mixed_ownership + + planner.stage_confidence(buffer, worker.confidence_stamp_buffer()) + stage_mixed_ownership(planner, worker, int(self.iter_counter), + gen_requests) + else: + planner.stage_confidence(buffer) + + def _dspark_confidence_draft_len(self, + scheduled_batch: ScheduledRequests) -> int: + """Agree one DSpark verify shape and stage confidence asynchronously.""" + max_draft_len = self.model_engine.max_draft_len + worker = self.model_engine._get_spec_worker() + planner = getattr(worker, "verify_planner", None) + if planner is None: + return max_draft_len + mixed = self.model_engine._dspark_mixed_runtime + if mixed is not None: + mixed.reset_step() + + runtime_cost_table = getattr(self.model_engine, + "_dspark_sps_cost_table", None) + if runtime_cost_table is not None and planner.cost_table is not runtime_cost_table: + planner.install_runtime_cost_table(runtime_cost_table) + + gen_requests = scheduled_batch.generation_requests + is_distributed = self.dist is not None and self.dist.tp_size > 1 + compute_windows = bool( + getattr(self.model_engine, "_dspark_trims_submitted_tokens", False)) + + ragged_lens = None + device_budget = None + exact_local = None + exact_table = planner.exact_cost_table + if exact_table is None: + raise RuntimeError( + "DSpark confidence scheduling has no authenticated exact SPS table" + ) + exact_cells = tuple( + getattr(self.model_engine, "_dspark_exact_candidate_cells", ())) + if not exact_cells: + exact_cells = exact_table.candidate_cells() + exact_requests: Optional[List[LlmRequest]] = gen_requests + exact_zero_real_adp_dummy = False + exact_requests, exact_zero_real_adp_dummy = ( + _classify_dspark_exact_generation_rows(gen_requests)) + + # A stale capacity must never re-rank a later batch. + self.model_engine._dspark_device_budget = None + if compute_windows: + if scheduled_batch.context_requests: + pass + else: + if exact_requests is not None: + exact_rows = [ + worker.confidence_row_for(request.py_request_id) + for request in exact_requests + ] + exact_local = planner.prepare_exact_sps_decision( + num_gen_requests=len(exact_requests), rows=exact_rows) + + local_rows = (exact_local.num_requests if exact_local is not None else + len(ragged_lens) if ragged_lens is not None else 0) + local_tokens = (sum( + 1 + int(value) + for value in ragged_lens) if ragged_lens is not None else 0) + local_payload = [ + 1 if exact_local is not None or ragged_lens is not None else 0, + int(local_rows), + int(local_tokens), + 1 if scheduled_batch.can_run_cuda_graph else 0, + ] + + exact_yield_scale = 1_000_000 + exact_identity_words = (0, ) * 8 + exact_identity_words = tuple( + getattr( + self.model_engine, + "_dspark_exact_identity_words", + exact_table.collective_identity_words, + )) + native_yield_wire, compact_yields_wire = _encode_dspark_exact_expected_yields( + exact_local=exact_local, + num_cells=len(exact_cells), + yield_scale=exact_yield_scale, + ) + local_payload.extend([ + 1, + 1 if exact_local is not None else 0, + len(exact_cells), + *exact_identity_words, + native_yield_wire, + *compact_yields_wire, + ]) + + runner = self.model_engine.cuda_graph_runner + local_payload.extend([ + int(scheduled_batch.batch_size), + 1 if planner.max_verify_len in getattr( + runner, "padding_dummy_requests", {}) else 0, + 1 if planner.max_verify_len in getattr( + runner, "secondary_padding_dummy_requests", {}) else 0, + ]) + if mixed is not None: + local_payload.extend( + mixed.local_vote( + scheduled_batch, + worker, + planner, + int(self.iter_counter), + int(self.dist.tp_rank) if is_distributed else 0, + executor_eligible=bool( + compute_windows + and self._dspark_adp_shape_cache_enabled()), + )) + payloads = (self.dist.tp_allgather_int64(local_payload).tolist() + if is_distributed else [local_payload]) + + exact_payload_end = _DSPARK_EXACT_WIRE_PREFIX_LEN + len(exact_cells) + expected_payload_len = exact_payload_end + _DSPARK_ADP_WIRE_TRAILER_LEN + mixed_peers = None + if mixed is not None: + payloads, mixed_peers = mixed.split_votes(payloads, + expected_payload_len) + if any(len(payload) != expected_payload_len for payload in payloads): + raise RuntimeError( + "DSpark ADP agreement payload has an inconsistent shape") + + exact_table_flags = [int(payload[4]) for payload in payloads] + if len(set(exact_table_flags)) != 1: + raise RuntimeError( + "DSpark ranks disagree on whether an exact SPS table is installed" + ) + if exact_table_flags[0]: + exact_cell_counts = [int(payload[6]) for payload in payloads] + exact_identities = [ + tuple(int(value) for value in payload[7:15]) + for payload in payloads + ] + if (len(set(exact_cell_counts)) != 1 + or exact_cell_counts[0] != len(exact_cells) + or len(set(exact_identities)) != 1): + raise RuntimeError( + "DSpark exact SPS ranks disagree on the authenticated " + "(G,V), T(G,V), or minimum-gain table identity") + + all_can_graph = all(int(payload[3]) for payload in payloads) + peer_stats = None + if not all(int(payload[0]) for payload in payloads): + ragged_lens = None + exact_local = None + else: + peer_stats = [[ + int(payload[1]), + int(payload[2]), 1 if all_can_graph else 0 + ] for payload in payloads] + + batch_size_index = exact_payload_end + padding_dummy_index = exact_payload_end + 1 + secondary_padding_dummy_index = exact_payload_end + 2 + peer_batch_sizes = tuple( + int(payload[batch_size_index]) for payload in payloads) + widest_batch_size = max(peer_batch_sizes, default=0) + common_padded_batch_size = runner._round_up_batch_size_with_draft_len( + widest_batch_size, planner.max_verify_len) + padding_needed = any( + int(payload[batch_size_index]) != common_padded_batch_size + for payload in payloads) + padding_supported = bool( + not padding_needed + or (runner.padding_enabled + and common_padded_batch_size <= runner.config.batch_size)) + padding_ready = bool( + common_padded_batch_size > 0 and padding_supported and all( + int(payload[batch_size_index]) == common_padded_batch_size + or int(payload[padding_dummy_index]) for payload in payloads)) + exact_policy_agreed = bool(all(int(payload[5]) for payload in payloads)) + agreement = ADPShapeAgreement( + iteration=int(self.iter_counter), + batch_identity=id(scheduled_batch), + local_batch_size=int(scheduled_batch.batch_size), + peer_batch_sizes=peer_batch_sizes, + all_can_graph=all_can_graph, + widest_batch_size=widest_batch_size, + graph_batch_size=common_padded_batch_size, + draft_len=int(planner.max_verify_len), + padding_ready=padding_ready, + ) + runner.adp_shape_agreement = agreement + if (mixed_peers is not None and agreement.can_queue + and mixed.apply(scheduled_batch, worker, planner, mixed_peers, + int(self.iter_counter))): + # Common context-aware get_padding_params supplies the BREAKABLE + # body. A decode-only peer must not replay its local full graph. + runner.adp_shape_agreement = dataclasses.replace( + agreement, + all_can_graph=False, + ragged_bucket=None, + reuse_graph_shape=False, + ) + self._stage_dspark_confidence(worker, planner, gen_requests, mixed) + return int(planner.max_verify_len) + if not agreement.can_queue: + runner.agreed_ragged_bucket = None + runner.ragged_pad_verify_len = 0 + runner.ragged_zero_real_high_rows = 0 + for request in gen_requests: + request.py_verify_len = None + return int(planner.max_verify_len) + + exact_shape = None + if exact_local is not None: + if not all_can_graph: + exact_local = None + else: + if any(not int(payload[5]) for payload in payloads): + raise RuntimeError( + "DSpark exact SPS collective payloads disagree on the " + "measured (G,V) grid") + common_graph_batch_size = _dspark_exact_common_graph_batch_size( + payloads, runner._round_up_batch_size) + if common_graph_batch_size == 0: + exact_local = None + else: + from ..speculative.dspark_planner import \ + select_exact_sps_candidate + + ( + native_expected_yield, + compact_expected_yields, + compact_max_yield_losses_per_request, + ) = _decode_dspark_exact_expected_yields( + payloads=payloads, + exact_cells=exact_cells, + graph_batch_size=common_graph_batch_size, + yield_scale=exact_yield_scale, + ) + selected_verifier_budget = select_exact_sps_candidate( + graph_batch_size=common_graph_batch_size, + native_expected_yield=native_expected_yield, + compact_expected_yields=compact_expected_yields, + compact_max_yield_losses_per_request=( + compact_max_yield_losses_per_request), + cost_table=exact_table, + ) + if selected_verifier_budget == 0 or not ( + _dspark_exact_secondary_padding_ready( + payloads, + graph_batch_size=common_graph_batch_size, + verifier_budget=selected_verifier_budget, + secondary_ready_index= + secondary_padding_dummy_index, + )): + exact_local = None + else: + ragged_lens, local_exact_budget, exact_pad_tokens = ( + planner.allocate_exact_sps_candidate( + exact_local, + graph_batch_size=common_graph_batch_size, + verifier_budget=selected_verifier_budget, + )) + if exact_local.num_requests == 0: + if not exact_zero_real_adp_dummy: + raise RuntimeError( + "zero-real exact SPS allocation has no " + "authenticated attention-DP dummy row") + _, remainder = divmod(selected_verifier_budget, + common_graph_batch_size) + scheduled_tokens = int(exact_pad_tokens) + int( + remainder > 0) + ragged_lens = [scheduled_tokens - 1] + exact_shape = ( + common_graph_batch_size, + selected_verifier_budget, + exact_pad_tokens, + ) + if planner.device_windows and local_exact_budget > 0: + device_budget = int(local_exact_budget) + + from ..speculative.dspark_verify import \ + _exact_cell_geometry + + peer_stats = [] + for payload in payloads: + peer_rows = int(payload[1]) + peer_geometry = _exact_cell_geometry( + num_real=peer_rows, + graph_batch_size=common_graph_batch_size, + verifier_budget=selected_verifier_budget, + min_verify_len=planner.cfg.min_verify_len, + max_verify_len=planner.max_verify_len, + ) + if peer_geometry is None: + raise RuntimeError( + "Selected exact SPS cell is infeasible for " + f"peer rows={peer_rows}, " + f"G={common_graph_batch_size}, " + f"V={selected_verifier_budget}") + _, peer_real_tokens, _ = peer_geometry + peer_stats.append([ + peer_rows, peer_real_tokens, + 1 if all_can_graph else 0 + ]) + + bucket = None + ragged_active = False + if ragged_lens is not None: + if exact_shape is None: + raise RuntimeError( + "DSpark exact allocation produced no measured graph shape") + if all_can_graph: + bucket = self.model_engine.fit_ragged_verify_lens( + gen_requests, + ragged_lens, + peer_stats=peer_stats, + exact_shape=exact_shape, + exact_zero_real=exact_zero_real_adp_dummy, + ) + ragged_active = bucket is not None + if ragged_active and device_budget is not None: + self.model_engine._dspark_device_budget = int(device_budget) + else: + runner.agreed_ragged_bucket = None + runner.ragged_pad_verify_len = 0 + runner.ragged_zero_real_high_rows = 0 + for request, window in zip(gen_requests, ragged_lens): + request.py_verify_len = int(window) + ragged_active = True + + reuse_graph_shape = bool(exact_policy_agreed and runner.enabled + and all_can_graph and padding_ready + and not runner.config.use_mrope) + _validate_dspark_exact_bucket(exact_shape=exact_shape, bucket=bucket) + if (reuse_graph_shape and exact_shape is not None + and exact_shape[0] != common_padded_batch_size): + raise RuntimeError( + "DSpark exact policy selected a graph batch size that " + "disagrees with the cached ADP padding shape") + runner.adp_shape_agreement = dataclasses.replace( + agreement, + ragged_bucket=None if bucket is None else int(bucket), + reuse_graph_shape=reuse_graph_shape, + ) + + if not ragged_active: + for request in gen_requests: + request.py_verify_len = None + runner.agreed_ragged_bucket = None + runner.ragged_pad_verify_len = 0 + runner.ragged_zero_real_high_rows = 0 + + self._stage_dspark_confidence(worker, planner, gen_requests, mixed) + return int(planner.max_verify_len) + def _handle_dynamic_draft_len(self, scheduled_batch: ScheduledRequests) -> None: - """Synchronize draft length and buffers before preparing resources.""" - update_draft_len(self.model_engine, - scheduled_batch, - speculation_permanently_disabled=self. - speculation_permanently_disabled) + """Synchronize draft width and confidence policy before KV admission.""" + spec_config = getattr(self.model_engine, "spec_config", None) + confidence_driven = bool( + spec_config is not None + and getattr(spec_config, "enable_confidence_scheduling", False)) + if not confidence_driven or self.speculation_permanently_disabled: + # Main owns the static, scheduled and disabled draft-buffer rules. + update_draft_len(self.model_engine, + scheduled_batch, + speculation_permanently_disabled=self. + speculation_permanently_disabled) + return + + signature = ( + int(self.iter_counter), + id(scheduled_batch), + int(scheduled_batch.batch_size), + int(scheduled_batch.num_context_requests), + int(scheduled_batch.num_generation_requests), + ) + if getattr(self, "_dspark_dynamic_handled_signature", + None) == signature: + return + runtime_draft_len = self._dspark_confidence_draft_len(scheduled_batch) + agreement = self.model_engine.cuda_graph_runner.adp_shape_agreement + if (agreement is not None and agreement.matches( + scheduled_batch, self.iter_counter, padded=False) + and not agreement.can_queue): + self.model_engine.runtime_draft_len = runtime_draft_len + else: + update_draft_len(self.model_engine, + scheduled_batch, + draft_len=runtime_draft_len) + self._dspark_dynamic_handled_signature = signature + + def _dspark_adp_shape_cache_enabled(self) -> bool: + runner = getattr(self.model_engine, "cuda_graph_runner", None) + return bool( + self.enable_attention_dp and self.dist is not None + and self.dist.tp_size > 1 and runner is not None + and getattr(self.model_engine, "_dspark_confidence_enabled", False) + and getattr(self, "kv_connector_manager", None) is None + and getattr(self, "kv_cache_transceiver", None) is None + and getattr(self, "drafter", None) is None + and getattr(self, "_pp_rebalance_drain_iters", None) is None) def _initialize_remote_tail_adp(self) -> None: """Reserve the execution-only idle row after final KV/graph warmup.""" @@ -4013,6 +4565,28 @@ def _can_queue(self, scheduled_batch): # can_queue_this_rank is for case that the batch is not empty on this rank, but empty on other ranks # For bs == 1, we cannot pad dummy request to make the batch non-empty since it will cause the batch size to be 2. # 1 for dummy request, 1 for the yet-to-complete but not-yet-updated request. + runner = getattr(self.model_engine, "cuda_graph_runner", None) + if self._dspark_adp_shape_cache_enabled(): + signature = ( + int(self.iter_counter), + id(scheduled_batch), + int(scheduled_batch.batch_size), + int(scheduled_batch.num_context_requests), + int(scheduled_batch.num_generation_requests), + ) + if getattr(self, "_dspark_dynamic_handled_signature", + None) != signature: + runner.adp_shape_agreement = None + self._handle_dynamic_draft_len(scheduled_batch) + agreement = runner.adp_shape_agreement + if agreement is not None and agreement.matches( + scheduled_batch, self.iter_counter, padded=False): + return agreement.can_queue, agreement.can_queue_this_rank + elif runner is not None: + # Optional surfaces can mutate the request set after policy + # selection, so never reuse their preceding agreement. + runner.adp_shape_agreement = None + if getattr(self, "_enable_remote_tail_adp", False): # An idle execution row needs no active-request or sampler slot. # Still veto an all-empty fleet (e.g. cancellations after the vote). @@ -5612,7 +6186,7 @@ def _maybe_finish_pp_rebalance(self) -> None: self._pp_rebalance_drain_iters = None def _suspend_padding_dummies_for_rebalance( - self, mgr: KVCacheManagerV2) -> List[Tuple[int, LlmRequest]]: + self, mgr: KVCacheManagerV2) -> List[Tuple[int, int, LlmRequest]]: """Suspend the CUDA-graph padding dummies before ``adjust()``. ``CUDAGraphRunner`` retains one padding dummy per captured draft @@ -5629,21 +6203,27 @@ def _suspend_padding_dummies_for_rebalance( and silently drop padded batches to eager mode for the rest of the process lifetime. - Returns the ``(draft_len, dummy)`` pairs that were suspended. + Returns ``(draft_len, variant, dummy)`` triples. Variant zero is the + low-row dummy and variant one is the high-row dummy. """ runner = getattr(self.model_engine, "cuda_graph_runner", None) if runner is None: return [] - suspended: List[Tuple[int, LlmRequest]] = [] - for draft_len, dummy in runner.padding_dummy_requests.items(): - if mgr.is_request_active(dummy.py_request_id): - mgr.suspend_request(dummy) - suspended.append((draft_len, dummy)) + suspended: List[Tuple[int, int, LlmRequest]] = [] + caches = ( + getattr(runner, "padding_dummy_requests", {}), + getattr(runner, "secondary_padding_dummy_requests", {}), + ) + for variant, cache in enumerate(caches): + for draft_len, dummy in cache.items(): + if mgr.is_request_active(dummy.py_request_id): + mgr.suspend_request(dummy) + suspended.append((draft_len, variant, dummy)) return suspended def _resume_padding_dummies_after_rebalance( self, mgr: KVCacheManagerV2, - suspended: List[Tuple[int, LlmRequest]]) -> None: + suspended: List[Tuple[int, int, LlmRequest]]) -> None: """Resume the padding dummies suspended for ``adjust()``. A real request that fails to resume is left suspended for the @@ -5661,7 +6241,8 @@ def _resume_padding_dummies_after_rebalance( runner = getattr(self.model_engine, "cuda_graph_runner", None) if runner is None: return - for draft_len, dummy in suspended: + failed_draft_lens = set() + for draft_len, variant, dummy in suspended: if mgr.resume_request(dummy): continue if getattr(self, "_enable_remote_tail_adp", @@ -5671,8 +6252,13 @@ def _resume_padding_dummies_after_rebalance( continue logger.warning( "Could not resume the CUDA graph padding dummy request " - f"(draft_len={draft_len}) after a KV pool rebalance; " + f"(draft_len={draft_len}, variant={variant}) after a KV pool " + "rebalance; " "releasing it and falling back to lazy re-creation.") + failed_draft_lens.add(draft_len) + # Resume both siblings before releasing a low/high pair. Releasing + # inline could free a resumed sibling and later try to resume it. + for draft_len in sorted(failed_draft_lens): runner.release_padding_dummy(self.resource_manager, draft_len) @contextmanager @@ -6356,6 +6942,10 @@ def _configure_csa2_remote_tail(self, request: LlmRequest) -> None: request.py_csa2_remote_tail_split = split def _validate_request(self, request: LlmRequest): + if getattr(self.model_engine, "_dspark_confidence_enabled", False): + from ..speculative.dspark import validate_confidence_request + + validate_confidence_request(self.model_engine.spec_config, request) self._configure_csa2_remote_tail(request) if (isinstance(self.kv_cache_manager, KVCacheManagerV2) and request.kv_cache_retention_config is not None): @@ -9048,7 +9638,7 @@ def _enqueue_responses(self, responses: Iterable[Tuple[int, LlmResponse]]): # (TODO: joyang) There are other types of responses, we need to sort out. if type( resp - ) == LlmResponse and req_id in self.result_wait_queues and self.result_wait_queues[ + ) is LlmResponse and req_id in self.result_wait_queues and self.result_wait_queues[ req_id] is not None: self.result_wait_queues[req_id].put_response.remote( resp.client_id, resp) diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 3e8690d8859e..4c740b7b831b 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -20,6 +20,7 @@ # ``DSv4DSparkDraftModel.forward`` rather than via mask-token cross-attention. import json +import math from collections import deque from dataclasses import dataclass, field from typing import TYPE_CHECKING, List, Optional @@ -37,6 +38,7 @@ if TYPE_CHECKING: from ...llmapi.llm_args import DSparkDecodingConfig + from ..pyexecutor.llm_request import LlmRequest # Unscored-slot fill: sigmoid saturates to 1.0, so an unwritten row schedules # as "verify the whole block" (fail-safe). @@ -60,7 +62,10 @@ def load_sts_temperatures_from_path(path: str) -> list[float]: ) if not isinstance(temps, list) or not temps: raise ValueError(f"{path} must contain a non-empty temperature list") - return [float(t) for t in temps] + temperatures = [float(t) for t in temps] + if any(not math.isfinite(value) or value <= 0.0 for value in temperatures): + raise ValueError("STS temperatures must be a finite positive numeric vector") + return temperatures def resolve_confidence_head(draft_model: torch.nn.Module) -> Optional[torch.nn.Module]: @@ -72,6 +77,35 @@ def resolve_confidence_head(draft_model: torch.nn.Module) -> Optional[torch.nn.M return getattr(stages[-1], "confidence_head", None) +def validate_confidence_request(spec_config: "DSparkDecodingConfig", request: "LlmRequest") -> None: + """Reject unsupported request controls before confidence-policy admission.""" + if not spec_config.enable_confidence_scheduling: + return + if getattr(request, "py_logits_post_processors", None): + raise ValueError( + "DSpark confidence scheduling does not support per-request logits post-processors" + ) + sampling = request.sampling_config + if not spec_config.use_rejection_sampling or sampling is None: + return + from tensorrt_llm.sampling_params import SamplingParams + + # Match SpecMetadata's actual draft-sampler classification; temperature + # alone is insufficient (for example, top_k=1 remains greedy). + greedy = SamplingParams.params_imply_greedy_decoding( + temperature=sampling.temperature, + top_k=sampling.top_k, + top_p=sampling.top_p, + min_p=sampling.min_p, + use_beam_search=False, + ) + if not greedy: + raise ValueError( + "DSpark confidence scheduling does not support advanced rejection draft sampling: " + "confidence is conditioned on the greedy proposal prefix" + ) + + def _dspark_position_ceiling(max_ctx: int, block_size: int, max_draft_len: int) -> int: """Return the number of RoPE entries needed by the DSv4 block drafter. @@ -233,6 +267,7 @@ def prepare(self): for rid in list(worker._req_to_slot.keys()): if rid not in current: slot = worker._req_to_slot.pop(rid) + worker._confidence_incarnations.pop(rid, None) worker._ctx_len[slot] = 0 worker._valid_len[slot] = 0 worker._position_initialized[slot] = False @@ -397,6 +432,13 @@ def __init__( # gen path (set on the host in prepare(), so the captured forward indexes # the rolling windows through a tensor instead of a python dict lookup). self._req_to_slot = {} # request_id -> slot index + # Enabled only for admitted host-mixed scheduling. Slot reuse cannot + # authenticate a different request incarnation with an older snapshot. + self._mixed_confidence_tracking = bool( + getattr(spec_config, "confidence_mixed_evidence", None) is not None + ) + self._confidence_incarnations: dict[int, int] = {} + self._confidence_incarnation_sequence = 0 self._free_slots = deque() # available slot indices self._batch_to_slot: Optional[torch.Tensor] = None # [max_batch] long, cuda # Index of the throwaway "scratch" window row that absorbs padded / @@ -451,7 +493,15 @@ def bump_draft_seq(self) -> int: sequence of the last completed draft. """ prev = self._draft_seq_host - self._draft_seq_host += 1 + self._draft_seq_host = (prev + 1) & ((1 << 31) - 1) + if self._draft_seq_host == 0: + # Recycled stamps must not authenticate an ancient score. The + # wrap step sees expected=-1 and uses neutral confidence; later + # steps can consume only newly produced scores or neutral rows. + if self._confidence_stamp is not None: + self._confidence_stamp.zero_() + if self._confidence_logits is not None: + self._confidence_logits.fill_(_NEUTRAL_CONFIDENCE_LOGIT) if self._draft_seq_cuda is not None: self._draft_seq_cuda.fill_(self._draft_seq_host) return prev @@ -699,6 +749,7 @@ def _assign_slot(self, req_id: int, reset: bool) -> int: """Get (or refresh) the slot for a request; reset clears its window.""" if reset and req_id in self._req_to_slot: old = self._req_to_slot.pop(req_id) + self._confidence_incarnations.pop(req_id, None) self._ctx_len[old] = 0 self._valid_len[old] = 0 self._position_initialized[old] = False @@ -712,6 +763,9 @@ def _assign_slot(self, req_id: int, reset: bool) -> int: ) slot = self._free_slots.popleft() self._req_to_slot[req_id] = slot + if self._mixed_confidence_tracking: + self._confidence_incarnation_sequence += 1 + self._confidence_incarnations[req_id] = self._confidence_incarnation_sequence self._ctx_len[slot] = 0 self._valid_len[slot] = 0 self._position_initialized[slot] = False @@ -723,6 +777,10 @@ def _assign_slot(self, req_id: int, reset: bool) -> int: self._confidence_logits[slot].fill_(_NEUTRAL_CONFIDENCE_LOGIT) return self._req_to_slot[req_id] + def confidence_incarnation_for(self, req_id: int) -> int | None: + """Host ownership generation, independent of a recycled device slot.""" + return self._confidence_incarnations.get(req_id) + def _seed_context_windows( self, draft_model, diff --git a/tensorrt_llm/_torch/speculative/dspark_mixed_evidence.py b/tensorrt_llm/_torch/speculative/dspark_mixed_evidence.py new file mode 100644 index 000000000000..bd6ba15dc7c4 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_mixed_evidence.py @@ -0,0 +1,514 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pre-construction evidence binding for measured BREAKABLE mixed costs. + +The optional mixed configuration binds measured costs and an admission receipt +to an independently observed live fingerprint. The runtime separately checks +the active route, request ownership and geometry before publishing row lengths. +Parsing authenticates evidence bytes and checks their claims; it does not run +the qualification performed by the evidence producer. +""" + +import hashlib +import json +from dataclasses import dataclass +from pathlib import Path + +from . import dspark_planner + +_BODY_BUCKETS = (96, 128, 192, 256, 384, 512) +_ACTION = "current-host-confidence-prefix-allocation-under-shared-body-capacity-v1" +_PRICE_SEMANTICS = "rank0_forward_pricing_proxy" +_GEOMETRY_FIELDS = ( + "context_rows", + "context_chunk_squared", + "context_position_product", + "context_compressed_product", + "generation_history_sum", + "generation_history_max", + "bootstrap_generations", + "real_context_requests", + "eligible_generations", + "context_requests", + "generation_requests", + "fixed_generation_rows", +) +_MAX_BYTES = 16 * 1024 * 1024 +_MAX_CELLS = 256 + + +@dataclass(frozen=True) +class MixedRuntimeIdentity: + """Independently observed live identity, not copied from the cost artifact. + + Support is limited to K3 and the eight-rank BREAKABLE body grid. The runner + checks embedded-model, PP/CP/attention-DP and capture readiness separately; + these properties are not inferred from a cost file. + """ + + source_tree_sha256: str + runtime_snapshot_sha256: str + model_config_sha256: str + gpu_class_sha256: str + topology_sha256: str + selector_identity_sha256: str + physical_k: int = 3 + world_size: int = 8 + body_buckets: tuple[int, ...] = _BODY_BUCKETS + execution_mode: str = "BREAKABLE" + + def __post_init__(self) -> None: + for field in ( + "source_tree_sha256", + "runtime_snapshot_sha256", + "model_config_sha256", + "gpu_class_sha256", + "topology_sha256", + "selector_identity_sha256", + ): + dspark_planner._require_sha256(getattr(self, field), field=field) + if ( + type(self.physical_k) is not int + or self.physical_k != 3 + or type(self.world_size) is not int + or self.world_size != 8 + or type(self.body_buckets) is not tuple + or any(type(bucket) is not int for bucket in self.body_buckets) + or self.body_buckets != _BODY_BUCKETS + or self.execution_mode != "BREAKABLE" + ): + raise ValueError("unsupported mixed physical K, topology or BREAKABLE grid") + + def payload(self) -> dict[str, object]: + return { + **{ + field: getattr(self, field) + for field in ( + "source_tree_sha256", + "runtime_snapshot_sha256", + "model_config_sha256", + "gpu_class_sha256", + "topology_sha256", + "selector_identity_sha256", + "physical_k", + "world_size", + "execution_mode", + ) + }, + "body_buckets": list(self.body_buckets), + } + + +@dataclass(frozen=True) +class MixedGeometry: + """One exact ordered-rank host-geometry/action key; never a decode ``(G,V)``. + + The twelve rank features match the accepted host price support dimensions. + They summarize host geometry, not effective device KV lengths or per-request + ownership. Ownership/incarnation/tickets must be checked separately at publish. + """ + + rank_features: tuple[tuple[int, ...], ...] + capacity: int + actual_rank_rows: tuple[int, ...] + native_rank_rows: tuple[int, ...] + + def __post_init__(self) -> None: + if ( + type(self.rank_features) is not tuple + or len(self.rank_features) != 8 + or type(self.actual_rank_rows) is not tuple + or type(self.native_rank_rows) is not tuple + or len(self.actual_rank_rows) != 8 + or len(self.native_rank_rows) != 8 + or type(self.capacity) is not int + or self.capacity not in _BODY_BUCKETS + ): + raise ValueError("mixed geometry requires eight ordered ranks and a captured capacity") + for rank, features in enumerate(self.rank_features): + if type(features) is not tuple or len(features) != len(_GEOMETRY_FIELDS): + raise ValueError("mixed rank feature dimensions differ") + for value in (*features, self.actual_rank_rows[rank], self.native_rank_rows[rank]): + dspark_planner._require_json_int(value, field="mixed geometry value", minimum=0) + if value >= 2**63: + raise ValueError("mixed geometry exceeds the existing wire integer bound") + ( + context, + _, + _, + _, + history_sum, + history_max, + bootstrap, + real_context, + eligible, + contexts, + generations, + fixed, + ) = features + if ( + real_context > contexts + or bootstrap > generations - eligible + or eligible > generations + or history_max > history_sum + or fixed != (generations - eligible) * 4 + or self.native_rank_rows[rank] != context + fixed + eligible * 4 + or not context + fixed + eligible + <= self.actual_rank_rows[rank] + <= min(self.native_rank_rows[rank], self.capacity) + ): + raise ValueError("mixed geometry does not reserve context and full-K fixed rows") + if ( + not any(row[7] for row in self.rank_features) + or self.actual_rank_rows == self.native_rank_rows + or _paid_bucket(self.native_rank_rows) is None + ): + raise ValueError("geometry is all-decode, noncompact or outside the captured grid") + + def payload(self) -> dict[str, object]: + return { + "rank_features": [dict(zip(_GEOMETRY_FIELDS, row)) for row in self.rank_features], + "capacity": self.capacity, + "actual_rank_rows": list(self.actual_rank_rows), + "native_rank_rows": list(self.native_rank_rows), + } + + +@dataclass(frozen=True) +class MixedEvidenceBinding: + """Exact artifact pins supplied by the optional mixed configuration.""" + + costs_sha256: str + admission_sha256: str + protocol_sha256: str + + def __post_init__(self) -> None: + for field in ("costs_sha256", "admission_sha256", "protocol_sha256"): + dspark_planner._require_sha256(getattr(self, field), field=field) + + +@dataclass(frozen=True) +class MixedCostCell: + geometry: MixedGeometry + native_lower_ms: float + candidate_upper_ms: float + measurement_result_sha256: str + + +@dataclass(frozen=True) +class AdmittedMixedCosts: + identity: MixedRuntimeIdentity + binding: MixedEvidenceBinding + minimum_predicted_gain: float + cells: tuple[MixedCostCell, ...] + + def lookup(self, geometry: MixedGeometry) -> MixedCostCell | None: + """Uncovered/invalid geometry stays native; no nearest-cell selection.""" + if type(geometry) is not MixedGeometry: + return None + return next((cell for cell in self.cells if cell.geometry == geometry), None) + + +@dataclass(frozen=True) +class MixedEvidenceResolution: + costs: AdmittedMixedCosts | None + reason: str + + +def load_mixed_evidence(config) -> MixedEvidenceResolution: + """Read independently pinned deployment artifacts before model execution. + + The live fingerprint is a separately produced artifact, never the table's + identity. Runtime route checks are additional to this byte binding. Invalid + optional evidence leaves mixed verification native, not confidence-off. + """ + if config is None: + return MixedEvidenceResolution(None, "mixed_evidence_absent_native") + try: + binding = MixedEvidenceBinding( + config.costs_sha256, config.admission_sha256, config.protocol_sha256 + ) + + def bounded_file(path: str) -> bytes: + with Path(path).open("rb") as handle: + value = handle.read(_MAX_BYTES + 1) + if not 0 < len(value) <= _MAX_BYTES: + raise ValueError("Mixed artifact exceeds the bounded byte contract") + return value + + live = dspark_planner._validate_exact_fields( + _read_pinned( + bounded_file(config.live_fingerprint_path), + config.live_fingerprint_sha256, + name="mixed live fingerprint", + ), + name="mixed live fingerprint", + fields={"schema", "identity"}, + ) + if live["schema"] != "dspark-mixed-live-fingerprint-v1": + raise ValueError("Unknown independently observed mixed fingerprint") + identity = dspark_planner._validate_exact_fields( + live["identity"], + name="mixed live identity", + fields={ + "source_tree_sha256", + "runtime_snapshot_sha256", + "model_config_sha256", + "gpu_class_sha256", + "topology_sha256", + "selector_identity_sha256", + "physical_k", + "world_size", + "body_buckets", + "execution_mode", + }, + ) + if type(identity["body_buckets"]) is not list: + raise ValueError("Mixed fingerprint body buckets must be a list") + observed = MixedRuntimeIdentity( + **{**identity, "body_buckets": tuple(identity["body_buckets"])} + ) + return resolve_mixed_evidence( + binding, + costs_raw=bounded_file(config.costs_path), + admission_raw=bounded_file(config.admission_path), + live_identity=observed, + ) + except ( + OSError, + TypeError, + ValueError, + KeyError, + AttributeError, + UnicodeError, + OverflowError, + RecursionError, + ) as exc: + return MixedEvidenceResolution( + None, f"mixed_evidence_invalid_native:{type(exc).__name__}:{exc}" + ) + + +def _paid_bucket(rows: tuple[int, ...]) -> int | None: + return next((bucket for bucket in _BODY_BUCKETS if bucket >= max(rows)), None) + + +def _read_pinned(raw: bytes, pin: str, *, name: str) -> dict[str, object]: + if type(raw) is not bytes or not 0 < len(raw) <= _MAX_BYTES: + raise ValueError(f"{name} must be bounded nonempty bytes") + if hashlib.sha256(raw).hexdigest() != pin: + raise ValueError(f"{name} byte digest differs") + + def unique(pairs: list[tuple[str, object]]) -> dict[str, object]: + result = {} + for key, value in pairs: + if key in result: + raise ValueError(f"{name} has duplicate fields") + result[key] = value + return result + + def reject_constant(value: str) -> None: + raise ValueError(f"{name} has nonfinite JSON: {value}") + + value = json.loads(raw, object_pairs_hook=unique, parse_constant=reject_constant) + if not isinstance(value, dict): + raise ValueError(f"{name} must be a JSON object") + return value + + +def _parse_geometry(value: object) -> MixedGeometry: + row = dspark_planner._validate_exact_fields( + value, + name="mixed geometry", + fields={ + "rank_features", + "capacity", + "actual_rank_rows", + "native_rank_rows", + }, + ) + if ( + type(row["rank_features"]) is not list + or type(row["actual_rank_rows"]) is not list + or type(row["native_rank_rows"]) is not list + ): + raise ValueError("mixed geometry arrays differ") + features = tuple( + tuple( + dspark_planner._validate_exact_fields( + rank, name="mixed rank", fields=set(_GEOMETRY_FIELDS) + )[field] + for field in _GEOMETRY_FIELDS + ) + for rank in row["rank_features"] + ) + return MixedGeometry( + features, row["capacity"], tuple(row["actual_rank_rows"]), tuple(row["native_rank_rows"]) + ) + + +def _validate( + binding: MixedEvidenceBinding, + costs_raw: bytes, + admission_raw: bytes, + identity: MixedRuntimeIdentity, +) -> AdmittedMixedCosts: + table = dspark_planner._validate_exact_fields( + _read_pinned(costs_raw, binding.costs_sha256, name="mixed costs"), + name="mixed costs", + fields={ + "schema", + "identity", + "protocol_sha256", + "action_semantics", + "pricing_semantics", + "synthetic", + "provenance", + "minimum_predicted_gain", + "cells", + }, + ) + receipt = dspark_planner._validate_exact_fields( + _read_pinned(admission_raw, binding.admission_sha256, name="mixed admission"), + name="mixed admission", + fields={ + "schema", + "identity", + "costs_sha256", + "protocol_sha256", + "action_semantics", + "pricing_semantics", + "production_admission", + "source_and_measurements_verified", + "accuracy_pass", + "runtime_invariants_pass", + "matched_workload_pass", + "accuracy_result_sha256", + "runtime_invariants_result_sha256", + "matched_workload_result_sha256", + "selector_replay_sha256", + "measurement_results_sha256", + "economics", + }, + ) + if ( + table["schema"] != "dspark-mixed-breakable-costs-v1" + or receipt["schema"] != "dspark-mixed-breakable-admission-v1" + or table["identity"] != identity.payload() + or receipt["identity"] != identity.payload() + or table["protocol_sha256"] != binding.protocol_sha256 + or receipt["protocol_sha256"] != binding.protocol_sha256 + or receipt["costs_sha256"] != binding.costs_sha256 + or any( + payload["action_semantics"] != _ACTION + or payload["pricing_semantics"] != _PRICE_SEMANTICS + for payload in (table, receipt) + ) + or table["synthetic"] is not False + or table["provenance"] != "closed_measured_breakable_mixed" + or any( + receipt[field] is not True + for field in ( + "production_admission", + "source_and_measurements_verified", + "accuracy_pass", + "runtime_invariants_pass", + "matched_workload_pass", + ) + ) + ): + raise ValueError("mixed source, execution, action or production evidence differs") + margin = dspark_planner._require_nonnegative_finite_number( + table["minimum_predicted_gain"], field="mixed gain margin" + ) + if margin < 0.01: + raise ValueError("mixed evidence cannot reduce the incumbent 1% immediate-goodput gate") + for field in ( + "accuracy_result_sha256", + "runtime_invariants_result_sha256", + "matched_workload_result_sha256", + "selector_replay_sha256", + "measurement_results_sha256", + ): + dspark_planner._require_sha256(receipt[field], field=field) + economics = dspark_planner._validate_exact_fields( + receipt["economics"], + name="mixed economics", + fields={ + "policy_steps", + "compact_choices", + "gross_compact_value_ms_lower_bound", + "fixed_confidence_path_overhead_ms_upper_bound", + "safety_margin_ms", + }, + ) + admitted, _ = dspark_planner.evaluate_confidence_workload_admission(**economics) + if not admitted: + raise ValueError("mixed workload economics are not positive") + if type(table["cells"]) is not list or not 0 < len(table["cells"]) <= _MAX_CELLS: + raise ValueError("mixed cells must be a bounded nonempty list") + cells = [] + for item in table["cells"]: + row = dspark_planner._validate_exact_fields( + item, + name="mixed cell", + fields={ + "geometry", + "native_lower_ms", + "candidate_upper_ms", + "measurement_result_sha256", + }, + ) + geometry = _parse_geometry(row["geometry"]) + if any(cell.geometry == geometry for cell in cells): + raise ValueError("duplicate mixed geometry/action cell") + cells.append( + MixedCostCell( + geometry, + dspark_planner._require_positive_finite_number( + row["native_lower_ms"], field="mixed native lower time" + ), + dspark_planner._require_positive_finite_number( + row["candidate_upper_ms"], field="mixed candidate upper time" + ), + dspark_planner._require_sha256( + row["measurement_result_sha256"], field="mixed measurement result" + ), + ) + ) + if ( + dspark_planner._canonical_json_sha256([cell.measurement_result_sha256 for cell in cells]) + != receipt["measurement_results_sha256"] + ): + raise ValueError("mixed admission does not bind all ordered measurement results") + return AdmittedMixedCosts(identity, binding, margin, tuple(cells)) + + +def resolve_mixed_evidence( + binding: MixedEvidenceBinding | None, + *, + costs_raw: bytes | None, + admission_raw: bytes | None, + live_identity: MixedRuntimeIdentity, +) -> MixedEvidenceResolution: + """Resolve once before construction; no decode object/flag is read or mutated. + + Invalid optional mixed evidence fails to native mixed verification. It must + never globally disable an already admitted decode policy, nor be upgraded by + copying a validation-only fingerprint or merely asserting a production flag. + """ + if binding is None: + return MixedEvidenceResolution(None, "mixed_evidence_absent_native") + try: + if ( + type(binding) is not MixedEvidenceBinding + or type(live_identity) is not MixedRuntimeIdentity + ): + raise ValueError("unknown mixed binding/live identity type") + return MixedEvidenceResolution( + _validate(binding, costs_raw, admission_raw, live_identity), "mixed_evidence_bound" + ) + except (TypeError, ValueError, KeyError, UnicodeError, OverflowError, RecursionError) as exc: + return MixedEvidenceResolution( + None, f"mixed_evidence_invalid_native:{type(exc).__name__}:{exc}" + ) diff --git a/tensorrt_llm/_torch/speculative/dspark_planner.py b/tensorrt_llm/_torch/speculative/dspark_planner.py index 812a6932b2db..0d87e590ab4b 100644 --- a/tensorrt_llm/_torch/speculative/dspark_planner.py +++ b/tensorrt_llm/_torch/speculative/dspark_planner.py @@ -175,9 +175,12 @@ class ExactSpsCostTable: max_draft_len: int minimum_predicted_gain: float = 0.01 iteration_drain_guard: ExactSpsDrainGuard | None = None + validate_artifact_freshness: bool = True identity_sha256: str = field(init=False) def __post_init__(self) -> None: + if type(self.validate_artifact_freshness) is not bool: + raise TypeError("validate_artifact_freshness must be a bool") max_draft_len = _require_exact_int(self.max_draft_len, field="max_draft_len", minimum=1) normalized: dict[int, ExactSpsCostRow] = {} for graph_batch_size, table in self.tables.items(): @@ -278,6 +281,9 @@ def __post_init__(self) -> None: for graph_batch_size, table in sorted(normalized.items()) ], } + if not self.validate_artifact_freshness: + # Reuse and strict policies must not share a collective identity. + identity_payload["validate_artifact_freshness"] = False object.__setattr__(self, "identity_sha256", _canonical_json_sha256(identity_payload)) @property @@ -371,36 +377,44 @@ def load_runtime_sps_cost_table( graph_batch_sizes: Sequence[int], max_draft_len: int, live_engine_fingerprint_path: str | Path | None = None, + validate_freshness: bool = True, ) -> tuple[ExactSpsCostTable, dict[str, object]]: - """Load an authenticated schema-v2 exact ``T(G,V)`` table. + """Load structurally checked exact costs, optionally checking freshness. - The source-head pin is intentional: a TRT-LLM version can cover source - changes that alter measured step costs. A changed source identity requires - requalification of the cost table and its workload-admission receipt. + Original artifact metadata is returned unchanged. Reuse does not assert + that recorded timings still describe the active source or GPU node. """ payload = _read_sps_cost_payload(path) schema_version = payload.get("schema_version") if type(schema_version) is not int or schema_version != 2: raise ValueError("Exact SPS cost artifacts require schema_version=2") - if live_engine_fingerprint_path is None: + if validate_freshness and live_engine_fingerprint_path is None: raise ValueError( "Schema-v2 exact SPS artifacts require an independently generated " "live engine fingerprint path" ) - live_engine_fingerprint = _read_sps_cost_payload(live_engine_fingerprint_path) + live_engine_fingerprint = ( + _read_sps_cost_payload(live_engine_fingerprint_path) + if live_engine_fingerprint_path is not None + else None + ) validate_sps_cost_table_payload( payload, graph_batch_sizes=graph_batch_sizes, max_draft_len=max_draft_len, live_engine_fingerprint=live_engine_fingerprint, + validate_freshness=validate_freshness, ) - return _build_exact_sps_cost_table(payload, max_draft_len=max_draft_len), payload + return _build_exact_sps_cost_table( + payload, max_draft_len=max_draft_len, validate_artifact_freshness=validate_freshness + ), payload def _build_exact_sps_cost_table( payload: dict[str, object], *, max_draft_len: int, + validate_artifact_freshness: bool = True, ) -> ExactSpsCostTable: exact_tables = payload["cost_tables"] assert isinstance(exact_tables, dict) @@ -452,6 +466,7 @@ def _build_exact_sps_cost_table( payload["minimum_predicted_gain"], field="minimum_predicted_gain" ), iteration_drain_guard=iteration_drain_guard, + validate_artifact_freshness=validate_artifact_freshness, ) @@ -698,6 +713,7 @@ def validate_sps_cost_table_payload( graph_batch_sizes: Sequence[int], max_draft_len: int, live_engine_fingerprint: dict[str, object] | None = None, + validate_freshness: bool = True, ) -> dict[str, object]: """Validate exact schema-v2 coverage and engine provenance. @@ -711,8 +727,9 @@ def validate_sps_cost_table_payload( graph_batch_sizes: Rank-local CUDA-graph row counts captured by the active engine. max_draft_len: Physical full draft-block length K. - live_engine_fingerprint: Independently supplied runtime fingerprint. - Validation refuses to use an artifact without this comparison. + live_engine_fingerprint: Optional independently supplied runtime identity. + validate_freshness: Require an exact live-identity comparison in addition + to mandatory schema, K, graph coverage and measurement provenance. Returns: The artifact's authenticated engine fingerprint. @@ -733,30 +750,37 @@ def validate_sps_cost_table_payload( assert isinstance(exact_tables, dict) assert isinstance(fingerprint, dict) - if live_engine_fingerprint is None: + if type(validate_freshness) is not bool: + raise TypeError("validate_freshness must be a bool") + if validate_freshness and live_engine_fingerprint is None: raise ValueError( "Schema-v2 SPS validation requires an independently supplied live engine fingerprint" ) - if live_engine_fingerprint is fingerprint: + if validate_freshness and live_engine_fingerprint is fingerprint: raise ValueError( "Live engine fingerprint must be supplied independently of the SPS artifact" ) - live_engine_fingerprint = _validate_exact_fields( - live_engine_fingerprint, - name="live engine fingerprint", - fields=_V2_FINGERPRINT_FIELDS, - ) - _validate_live_fingerprint_values(live_engine_fingerprint) + if live_engine_fingerprint is not None: + live_engine_fingerprint = _validate_exact_fields( + live_engine_fingerprint, + name="live engine fingerprint", + fields=_V2_FINGERPRINT_FIELDS, + ) + _validate_live_fingerprint_values(live_engine_fingerprint) if payload["engine_fingerprint_sha256"] != _canonical_json_sha256(fingerprint): raise ValueError("SPS engine_fingerprint SHA256 does not match payload") - mismatches = sorted( - key for key in _V2_FINGERPRINT_FIELDS if fingerprint[key] != live_engine_fingerprint[key] - ) - if mismatches: - raise ValueError( - "SPS engine_fingerprint does not match active runtime: " + ", ".join(mismatches) + if validate_freshness: + assert live_engine_fingerprint is not None + mismatches = sorted( + key + for key in _V2_FINGERPRINT_FIELDS + if fingerprint[key] != live_engine_fingerprint[key] ) + if mismatches: + raise ValueError( + "SPS engine_fingerprint does not match active runtime: " + ", ".join(mismatches) + ) configured_graph_batch_sizes = { _require_json_int(value, field="runtime graph batch size", minimum=1) @@ -1130,8 +1154,8 @@ def select_exact_sps_candidate( The caller supplies globally agreed expected yields after the phase-2 allgather. This pure layer only compares exact measured cells; it cannot - create a ragged budget or synchronize ranks. Compact serving fails closed - until the exact table carries matched iteration/drain metadata. Every + create a ragged budget or synchronize ranks. Strict validation retains + native verification without matched iteration/drain metadata. Every measured tier that clears the aggregate immediate-goodput guard is reranked by its guarded group-E2E value. Drain loss is the maximum expected loss per real request on any active attention-DP rank, derived from the same @@ -1199,10 +1223,10 @@ def select_exact_sps_candidate( return 0 drain_guard = cost_table.iteration_drain_guard if drain_guard is None: - # The matched mean-output statistic is workload-specific and cannot be - # inferred from T(G,V). An older exact artifact therefore remains a - # valid native-static control but is not allowed to enable compact V. - return 0 + if cost_table.validate_artifact_freshness: + return 0 + # Without optional drain evidence, reuse prices only immediate goodput. + return int(max(goodput_budgets, key=lambda budget: (compact_scores[budget], budget))) native_step_ms = cost_table.step_time(0, graph_batch_size) tail_step_ms = cost_table.step_time(0, drain_guard.tail_graph_batch_size) group_values_ms = {} diff --git a/tensorrt_llm/_torch/speculative/dspark_schedule.py b/tensorrt_llm/_torch/speculative/dspark_schedule.py index e9bd692b6259..2e49e64791b6 100644 --- a/tensorrt_llm/_torch/speculative/dspark_schedule.py +++ b/tensorrt_llm/_torch/speculative/dspark_schedule.py @@ -27,12 +27,16 @@ from dataclasses import dataclass import torch +import triton +import triton.language as tl __all__ = [ "DSparkScheduleConfig", + "DSparkFusedScheduleError", "HOST_POLICY_WINDOWS_SNAPSHOT_OUTPUT", "NATIVE_UNIFORM_VERIFY_OUTPUT", "compute_survival", + "schedule_verify_lens_topk_fused_fill", "schedule_verify_lens_topk", ] @@ -50,6 +54,15 @@ NATIVE_UNIFORM_VERIFY_OUTPUT = "native_uniform_verify" +class DSparkFusedScheduleError(RuntimeError): + """A synchronous optional fused-scheduler dispatch failure. + + Input and graph-contract errors remain ordinary ``ValueError`` exceptions; + callers may safely disable one fused G only for this wrapped launch error + and rerun the established tensor implementation. + """ + + @dataclass(frozen=True) class DSparkScheduleConfig: """Bounds for the per-request verify length. @@ -179,3 +192,269 @@ def schedule_verify_lens_topk( granted = selected.view(bs, schedulable).sum(dim=1).to(torch.int32) verify_lens += granted return torch.clamp(verify_lens, min=floor, max=cfg.resolved_max_verify_len) + + +@triton.jit(do_not_specialize=["num_real", "budget"]) +def _schedule_topk_rank_kernel( + survival_ptr, + output_ptr, + stride_row, + stride_position, + num_real, + budget, + survival_eps, + FLOOR: tl.constexpr, + COLS: tl.constexpr, + NUM_CANDIDATES: tl.constexpr, + BLOCK_COLS: tl.constexpr, + BLOCK_CP: tl.constexpr, +): + request = tl.program_id(0) + position = tl.arange(0, BLOCK_COLS) + candidate_mask = position < COLS + survival = tl.load( + survival_ptr + request * stride_row + (position + FLOOR).to(tl.int64) * stride_position, + mask=candidate_mask, + other=0.0, + ).to(tl.float32) + valid = candidate_mask & (survival >= survival_eps) + masked_survival = tl.where(valid, survival, float("-inf")) + + rank = tl.zeros([BLOCK_COLS], dtype=tl.int32) + for peer_start in range(0, NUM_CANDIDATES, BLOCK_CP): + peer = peer_start + tl.arange(0, BLOCK_CP) + peer_mask = peer < NUM_CANDIDATES + peer_request = peer // COLS + peer_position = peer % COLS + peer_survival = tl.load( + survival_ptr + + peer_request.to(tl.int64) * stride_row + + (peer_position + FLOOR).to(tl.int64) * stride_position, + mask=peer_mask, + other=0.0, + ).to(tl.float32) + peer_valid = peer_mask & (peer_request < num_real) & (peer_survival >= survival_eps) + peer_masked_survival = tl.where(peer_valid, peer_survival, float("-inf")) + + higher = peer_masked_survival[None, :] > masked_survival[:, None] + equal = peer_masked_survival[None, :] == masked_survival[:, None] + earlier = (peer_position[None, :] < position[:, None]) | ( + (peer_position[None, :] == position[:, None]) & (peer_request[None, :] < request) + ) + before = (higher | (equal & earlier)) & peer_valid[None, :] + rank += tl.sum(before.to(tl.int32), axis=1) + + selected = valid & (rank < budget) + tl.store(output_ptr + request, tl.sum(selected.to(tl.int32), axis=0)) + + +@triton.jit(do_not_specialize=["num_real", "budget", "pad_len", "graph_num_tokens"]) +def _schedule_topk_fill_kernel( + output_ptr, + num_rows, + num_real, + budget, + pad_len, + graph_num_tokens, + TOKEN_FLOOR: tl.constexpr, + MAX_TOKEN_LEN: tl.constexpr, + HAS_CANDIDATES: tl.constexpr, + BLOCK: tl.constexpr, +): + rows = tl.arange(0, BLOCK) + row_mask = rows < num_rows + is_real = row_mask & (rows < num_real) + selected_extra = tl.load( + output_ptr + rows, + mask=is_real & (budget > 0) & HAS_CANDIDATES, + other=0, + ).to(tl.int32) + token_len = tl.where(is_real, TOKEN_FLOOR + selected_extra, pad_len) + + # Reproduce fill_bucket_device's exact-pad real-row phase. The explicit + # policy budget controls only the confidence-ranked allocation. Any + # already-paid graph remainder is filled round-robin in request-index + # order; fusion must not promote it into another confidence budget. + spare = graph_num_tokens - tl.sum(tl.where(row_mask, token_len, 0), axis=0) + for _ in range(MAX_TOKEN_LEN - 1): + headroom = is_real & (token_len < MAX_TOKEN_LEN) + in_cycle = tl.cumsum(headroom.to(tl.int32), axis=0) <= spare + grant = headroom & in_cycle + token_len += grant.to(tl.int32) + spare -= tl.sum(grant.to(tl.int32), axis=0) + tl.store(output_ptr + rows, token_len, mask=row_mask) + + +def _schedule_verify_lens_topk_fused_fill_triton( + *, + survival: torch.Tensor, + budget: int, + num_real: int, + pad_len: int, + graph_num_tokens: int, + cfg: DSparkScheduleConfig, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Triton implementation of policy top-k plus exact round-robin fill.""" + num_rows = survival.shape[0] + if num_rows == 0: + return torch.empty(0, dtype=torch.int32, device=survival.device) if out is None else out[:0] + if num_rows > 256: + raise ValueError(f"fused top-k scheduler supports at most 256 rows, got {num_rows}") + schedulable = int(cfg.schedulable_per_request) + output = ( + torch.empty(num_rows, dtype=torch.int32, device=survival.device) + if out is None + else out[:num_rows] + ) + + # Specialize rank work by the captured graph row count, never by the + # step-varying real-row count. This keeps one compiled kernel per all-G + # graph key; peer rows past num_real are masked inside the kernel. + num_candidates = num_rows * schedulable + ranked = budget > 0 and num_real > 0 and num_candidates > 0 + if ranked: + block_cols = triton.next_power_of_2(schedulable) + block_cp = 128 + try: + _schedule_topk_rank_kernel[(num_real,)]( + survival, + output, + survival.stride(0), + survival.stride(1), + num_real, + budget, + float(cfg.survival_eps), + FLOOR=int(cfg.min_verify_len), + COLS=schedulable, + NUM_CANDIDATES=num_candidates, + BLOCK_COLS=block_cols, + BLOCK_CP=block_cp, + ) + except Exception as exc: + raise DSparkFusedScheduleError("DSpark fused top-k rank launch failed") from exc + finalize_block = triton.next_power_of_2(num_rows) + try: + _schedule_topk_fill_kernel[(1,)]( + output, + num_rows, + num_real, + budget, + pad_len, + graph_num_tokens, + TOKEN_FLOOR=int(cfg.min_verify_len) + 1, + MAX_TOKEN_LEN=int(cfg.resolved_max_verify_len) + 1, + HAS_CANDIDATES=schedulable > 0, + BLOCK=finalize_block, + ) + except Exception as exc: + raise DSparkFusedScheduleError("DSpark fused round-robin fill launch failed") from exc + return output + + +def schedule_verify_lens_topk_fused_fill( + *, + survival: torch.Tensor, + budget: int, + num_real: int, + pad_len: int, + cfg: DSparkScheduleConfig, + graph_num_tokens: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Fuse the established top-k schedule and exact round-robin bucket fill. + + The explicit ``budget`` retains exactly the semantics of + :func:`schedule_verify_lens_topk`. Pad rows receive exactly ``pad_len`` + tokens; any graph-token remainder is granted to real rows in index-order + round-robin cycles, exactly as ``fill_bucket_device(..., pad_fill=...)``. + CUDA uses two deterministic Triton launches. CPU executes the established + tensor operations as the functional oracle. When ``out`` is supplied, the + first ``num_rows`` int32 elements are overwritten and returned; CUDA then + performs no output allocation. + """ + if survival.dim() != 2: + raise ValueError(f"survival must be [bs, K], got shape {tuple(survival.shape)}") + num_rows, num_positions = survival.shape + if out is not None: + if out.device != survival.device: + raise ValueError("out must share survival.device") + if out.dtype != torch.int32: + raise TypeError(f"out must have dtype torch.int32, got {out.dtype}") + if out.dim() != 1 or out.numel() < num_rows: + raise ValueError(f"out must be a 1-D tensor with at least {num_rows} elements") + if not out.is_contiguous(): + raise ValueError("out must be contiguous") + if num_positions != cfg.block_size: + raise ValueError( + f"survival has {num_positions} positions but block_size is {cfg.block_size}" + ) + budget = int(budget) + num_real = int(num_real) + pad_len = int(pad_len) + if cfg.survival_eps <= 0.0: + raise ValueError( + "fused top-k scheduling requires survival_eps > 0 so zeroed pad " + "rows cannot compete for the legacy policy budget" + ) + if not 0 <= num_real <= num_rows: + raise ValueError(f"num_real must be in [0, {num_rows}], got {num_real}") + max_token_len = int(cfg.resolved_max_verify_len) + 1 + if not 1 <= pad_len <= max_token_len: + raise ValueError(f"pad_len must be in [1, {max_token_len}], got {pad_len}") + + graph_num_tokens = int(graph_num_tokens) + token_floor = int(cfg.min_verify_len) + 1 + pad_tokens = (num_rows - num_real) * pad_len + minimum_tokens = num_real * token_floor + pad_tokens + maximum_tokens = num_real * max_token_len + pad_tokens + scheduled_capacity = min(max(budget, 0), num_real * int(cfg.schedulable_per_request)) + if not minimum_tokens + scheduled_capacity <= graph_num_tokens <= maximum_tokens: + raise ValueError( + "fused schedule cannot realize the selected CUDA graph without " + "changing policy semantics: " + f"expected [{minimum_tokens + scheduled_capacity}, {maximum_tokens}] " + f"tokens, graph requires {graph_num_tokens}" + ) + + if survival.is_cuda: + return _schedule_verify_lens_topk_fused_fill_triton( + survival=survival, + budget=budget, + num_real=num_real, + pad_len=pad_len, + graph_num_tokens=graph_num_tokens, + cfg=cfg, + out=out, + ) + + rows = torch.arange(num_rows, device=survival.device) + real_survival = torch.where( + (rows < num_real).unsqueeze(1), + survival, + torch.full_like(survival, -1.0), + ) + scheduled = schedule_verify_lens_topk( + survival=real_survival, + budget=budget, + cfg=cfg, + ) + token_lens = torch.where( + rows < num_real, + scheduled + 1, + torch.full_like(scheduled, pad_len), + ) + spare = graph_num_tokens - int(token_lens.sum()) + for _ in range(max_token_len - 1): + if spare <= 0: + break + headroom = (rows < num_real) & (token_lens < max_token_len) + grant = headroom & (torch.cumsum(headroom.to(torch.int64), dim=0) <= spare) + token_lens += grant.to(torch.int32) + spare -= int(grant.sum()) + if spare != 0: + raise ValueError(f"fused schedule left {spare} graph tokens after round-robin fill") + if out is None: + return token_lens + out[:num_rows].copy_(token_lens) + return out[:num_rows] diff --git a/tensorrt_llm/_torch/speculative/dspark_select.py b/tensorrt_llm/_torch/speculative/dspark_select.py index 800a6e81ac8d..dff79b5f8454 100644 --- a/tensorrt_llm/_torch/speculative/dspark_select.py +++ b/tensorrt_llm/_torch/speculative/dspark_select.py @@ -18,9 +18,11 @@ 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. +per-token row maps is device-resident. The production exact-tier caller runs +it immediately before replay, after the previous step's draft wrote the +confidence buffer and before this step's verify attention reads the layout. +The tensor-control fallback remains suitable for capture by a caller whose +graph key owns all of its replay-varying controls. The split of responsibilities follows the paper's dual-timescale design: @@ -30,32 +32,47 @@ * 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. +Per-row staleness degrades to the neutral row (verify everything for that +request), mirroring the host planner's fail-open semantics. Exact-tier +feasibility is checked from host-known scalars before the fused launch; an +unsupported fused shape falls back to the established tensor schedule/fill. """ from dataclasses import dataclass -from typing import Callable, Optional +from typing import Callable, Optional, Union import torch +import triton +import triton.language as tl from .dspark_schedule import ( NEUTRAL_CONFIDENCE_LOGIT, + DSparkFusedScheduleError, DSparkScheduleConfig, compute_survival, schedule_verify_lens_topk, + schedule_verify_lens_topk_fused_fill, ) from .ragged_helpers import build_qo_indptr, build_row_maps_device, fill_bucket_device __all__ = [ "DeviceWindowResult", + "DeviceWindowWorkspace", + "DSparkCompactLayoutError", "gather_packed_draft_tokens", + "materialize_compact_layout", "select_windows_device", ] +class DSparkCompactLayoutError(DSparkFusedScheduleError): + """Synchronous optional-layout failure with snapshot recovery state.""" + + def __init__(self, message: str, *, past_seen_valid: bool) -> None: + super().__init__(message) + self.past_seen_valid = past_seen_valid + + @dataclass class DeviceWindowResult: """Everything the ragged layout consumers need, all device-resident. @@ -75,6 +92,221 @@ class DeviceWindowResult: qo_indptr: torch.Tensor req_idx: torch.Tensor kv_correction: torch.Tensor + workspace: Optional["DeviceWindowWorkspace"] = None + + +@dataclass +class DeviceWindowWorkspace: + """Stable-address scratch for the fused compact-layout prologue. + + A workspace is owned by one model engine and sliced to the active graph + shape. Keeping these tensors alive eliminates steady-state scheduler, + prefix-sum, row-map, and position-snapshot allocations. The tensor + fallback intentionally does not use this object, preserving it as an + independent oracle. + """ + + verify_lens: torch.Tensor + qo_indptr: torch.Tensor + req_idx: torch.Tensor + kv_correction: torch.Tensor + past_seen: torch.Tensor + + @classmethod + def allocate( + cls, + *, + max_rows: int, + max_tokens: int, + device: Union[str, torch.device], + ) -> "DeviceWindowWorkspace": + if max_rows < 0 or max_tokens < 0: + raise ValueError("workspace capacities must be non-negative") + return cls( + verify_lens=torch.empty(max_rows, dtype=torch.int32, device=device), + qo_indptr=torch.empty(max_rows + 1, dtype=torch.int32, device=device), + req_idx=torch.empty(max_tokens, dtype=torch.int64, device=device), + kv_correction=torch.empty(max_tokens, dtype=torch.int32, device=device), + past_seen=torch.empty(max_rows, dtype=torch.int32, device=device), + ) + + def validate(self, *, device: torch.device, num_rows: int, num_tokens: int) -> None: + tensors = ( + ("verify_lens", self.verify_lens, torch.int32, num_rows), + ("qo_indptr", self.qo_indptr, torch.int32, num_rows + 1), + ("req_idx", self.req_idx, torch.int64, num_tokens), + ("kv_correction", self.kv_correction, torch.int32, num_tokens), + ("past_seen", self.past_seen, torch.int32, num_rows), + ) + for name, tensor, dtype, capacity in tensors: + if tensor.device != device: + raise ValueError(f"workspace.{name} must be on {device}") + if tensor.dtype != dtype: + raise TypeError(f"workspace.{name} must have dtype {dtype}, got {tensor.dtype}") + if tensor.dim() != 1 or tensor.numel() < capacity: + raise ValueError( + f"workspace.{name} needs capacity {capacity}, got shape {tuple(tensor.shape)}" + ) + if not tensor.is_contiguous(): + raise ValueError(f"workspace.{name} must be contiguous") + + +@triton.jit +def _materialize_row_maps_kernel( + verify_lens_ptr, + qo_indptr_ptr, + req_idx_ptr, + correction_ptr, + MAX_TOKEN_LEN: tl.constexpr, +): + row = tl.program_id(0) + offsets = tl.arange(0, MAX_TOKEN_LEN) + length = tl.load(verify_lens_ptr + row) + token = tl.load(qo_indptr_ptr + row) + offsets + mask = offsets < length + tl.store(req_idx_ptr + token, row, mask=mask) + tl.store(correction_ptr + token, offsets - length + 1, mask=mask) + + +def _build_row_maps_workspace( + *, + workspace: DeviceWindowWorkspace, + num_rows: int, + graph_num_tokens: int, + max_token_len: int, +) -> None: + verify_lens = workspace.verify_lens[:num_rows] + qo_indptr = workspace.qo_indptr[: num_rows + 1] + qo_indptr[0].zero_() + if num_rows == 0: + return + torch.cumsum(verify_lens, dim=0, out=qo_indptr[1:]) + if verify_lens.is_cuda: + try: + _materialize_row_maps_kernel[(num_rows,)]( + verify_lens, + qo_indptr, + workspace.req_idx, + workspace.kv_correction, + MAX_TOKEN_LEN=triton.next_power_of_2(max_token_len), + ) + except Exception as exc: + raise DSparkFusedScheduleError("DSpark fused row-map materialization failed") from exc + return + + # CPU is a hardware-independent differential-test path. Production CUDA + # never enters the allocation-heavy tensor oracle below. + req_idx, correction = build_row_maps_device(verify_lens, graph_num_tokens=graph_num_tokens) + workspace.req_idx[:graph_num_tokens].copy_(req_idx) + workspace.kv_correction[:graph_num_tokens].copy_(correction) + + +@triton.jit(do_not_specialize=["num_rows"]) +def _snapshot_past_seen_kernel( + old_qo_ptr, + position_ids_ptr, + past_seen_ptr, + num_rows, + BLOCK: tl.constexpr, +): + rows = tl.arange(0, BLOCK) + mask = rows < num_rows + starts = tl.load(old_qo_ptr + rows, mask=mask, other=0) + past_seen = tl.load(position_ids_ptr + starts, mask=mask, other=0) + tl.store(past_seen_ptr + rows, past_seen, mask=mask) + + +@triton.jit( + do_not_specialize=[ + "new_tokens_stride_token", + "new_tokens_stride_slot", + "next_draft_stride_slot", + "next_draft_stride_token", + "num_tokens", + "num_real", + "real_tokens", + ] +) +def _materialize_compact_tokens_kernel( + req_idx_ptr, + qo_indptr_ptr, + past_seen_ptr, + batch_slots_ptr, + new_tokens_ptr, + new_tokens_lens_ptr, + next_draft_tokens_ptr, + input_ids_ptr, + position_ids_ptr, + previous_pos_indices_ptr, + previous_pos_id_offsets_ptr, + draft_tokens_ptr, + new_tokens_stride_token, + new_tokens_stride_slot, + next_draft_stride_slot, + next_draft_stride_token, + num_tokens, + num_real, + real_tokens, + BLOCK: tl.constexpr, +): + token = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + token_mask = token < num_tokens + row = tl.load(req_idx_ptr + token, mask=token_mask, other=0).to(tl.int32) + row_start = tl.load(qo_indptr_ptr + row, mask=token_mask, other=0) + offset = token - row_start + owner = tl.minimum(row, num_real - 1) + slot = tl.load(batch_slots_ptr + owner, mask=token_mask, other=0) + + new_token = tl.load( + new_tokens_ptr + offset * new_tokens_stride_token + slot * new_tokens_stride_slot, + mask=token_mask, + other=0, + ) + tl.store(input_ids_ptr + token, new_token, mask=token_mask) + past_seen = tl.load(past_seen_ptr + row, mask=token_mask, other=0) + tl.store(position_ids_ptr + token, past_seen + offset, mask=token_mask) + + real_mask = token < real_tokens + tl.store(previous_pos_indices_ptr + token, slot, mask=real_mask) + new_len = tl.load(new_tokens_lens_ptr + slot, mask=real_mask, other=0) + tl.store(previous_pos_id_offsets_ptr + token, new_len, mask=real_mask) + + draft_mask = token_mask & (row < num_real) & (offset > 0) + draft_token = tl.load( + next_draft_tokens_ptr + + slot * next_draft_stride_slot + + (offset - 1) * next_draft_stride_token, + mask=draft_mask, + other=0, + ) + draft_index = row_start - row + offset - 1 + tl.store(draft_tokens_ptr + draft_index, draft_token, mask=draft_mask) + + +@triton.jit(do_not_specialize=["num_rows"]) +def _commit_compact_rows_kernel( + new_lens_ptr, + new_qo_ptr, + old_lens_ptr, + old_qo_ptr, + kv_lens_ptr, + previous_kv_offsets_ptr, + num_rows, + BLOCK: tl.constexpr, +): + rows = tl.arange(0, BLOCK) + row_mask = rows < num_rows + new_lens = tl.load(new_lens_ptr + rows, mask=row_mask, other=0) + old_lens = tl.load(old_lens_ptr + rows, mask=row_mask, other=0) + delta = new_lens - old_lens + kv_lens = tl.load(kv_lens_ptr + rows, mask=row_mask, other=0) + previous_offsets = tl.load(previous_kv_offsets_ptr + rows, mask=row_mask, other=0) + tl.store(kv_lens_ptr + rows, kv_lens + 2 * delta, mask=row_mask) + tl.store(previous_kv_offsets_ptr + rows, previous_offsets - delta, mask=row_mask) + tl.store(old_lens_ptr + rows, new_lens, mask=row_mask) + qo_mask = rows <= num_rows + new_qo = tl.load(new_qo_ptr + rows, mask=qo_mask, other=0) + tl.store(old_qo_ptr + rows, new_qo, mask=qo_mask) def gather_packed_draft_tokens( @@ -113,18 +345,189 @@ def gather_packed_draft_tokens( return next_draft_tokens[slots, offsets] +def materialize_compact_layout( + *, + result: DeviceWindowResult, + old_verify_lens: torch.Tensor, + old_qo_indptr: torch.Tensor, + batch_slots: torch.Tensor, + new_tokens: torch.Tensor, + new_tokens_lens: torch.Tensor, + next_draft_tokens: torch.Tensor, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + previous_pos_indices: torch.Tensor, + previous_pos_id_offsets: torch.Tensor, + draft_tokens: torch.Tensor, + kv_lens: torch.Tensor, + previous_kv_lens_offsets: torch.Tensor, + num_real: int, + real_tokens: int, +) -> None: + """Materialize a selected compact layout without per-step temporaries. + + CUDA snapshots row positions, writes token/position/overlap/draft buffers, + and commits the row delta with three allocation-free Triton launches. A + CPU implementation mirrors the former tensor code for differential tests; + it is deliberately not a second optimized production path. + + The row commit is last. Therefore a synchronous failure in either earlier + launch leaves the staged lens/qo/KV state intact and lets the caller rerun + the established tensor fallback safely. + """ + workspace = result.workspace + if workspace is None: + raise ValueError("compact materialization requires a workspace-backed result") + num_rows = result.verify_lens.numel() + num_tokens = result.req_idx.numel() + num_real = int(num_real) + real_tokens = int(real_tokens) + if not 0 < num_real <= num_rows: + raise ValueError(f"num_real must be in [1, {num_rows}], got {num_real}") + if not 0 <= real_tokens <= num_tokens: + raise ValueError(f"real_tokens must be in [0, {num_tokens}], got {real_tokens}") + device = result.verify_lens.device + workspace.validate(device=device, num_rows=num_rows, num_tokens=num_tokens) + + one_dimensional = ( + ("old_verify_lens", old_verify_lens, num_rows), + ("old_qo_indptr", old_qo_indptr, num_rows + 1), + ("batch_slots", batch_slots, num_real), + ("new_tokens_lens", new_tokens_lens, 1), + ("input_ids", input_ids, num_tokens), + ("position_ids", position_ids, num_tokens), + ("previous_pos_indices", previous_pos_indices, real_tokens), + ("previous_pos_id_offsets", previous_pos_id_offsets, real_tokens), + ("draft_tokens", draft_tokens, real_tokens - num_real), + ("kv_lens", kv_lens, num_rows), + ("previous_kv_lens_offsets", previous_kv_lens_offsets, num_rows), + ) + for name, tensor, capacity in one_dimensional: + if tensor.device != device: + raise ValueError(f"{name} must be on {device}") + if tensor.dim() != 1 or tensor.numel() < capacity: + raise ValueError(f"{name} needs capacity {capacity}, got shape {tuple(tensor.shape)}") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + if new_tokens.device != device or new_tokens.dim() not in (2, 3): + raise ValueError("new_tokens must be a 2-D or 3-D tensor on the layout device") + if new_tokens.dim() == 3 and new_tokens.shape[2] != 1: + raise ValueError("new_tokens trailing beam dimension must be one") + if next_draft_tokens.device != device or next_draft_tokens.dim() != 2: + raise ValueError("next_draft_tokens must be a 2-D tensor on the layout device") + if old_verify_lens.dtype != torch.int32 or old_qo_indptr.dtype != torch.int32: + raise TypeError("old verify_lens and qo_indptr must have dtype torch.int32") + if position_ids.dtype != torch.int32 or workspace.past_seen.dtype != position_ids.dtype: + raise TypeError("position_ids and workspace.past_seen must have dtype torch.int32") + real_draft = real_tokens - num_real + if real_draft < 0: + raise ValueError("real_tokens must include one bonus token per real row") + + if result.verify_lens.is_cuda: + row_block = triton.next_power_of_2(num_rows + 1) + token_block = 256 + try: + _snapshot_past_seen_kernel[(1,)]( + old_qo_indptr, + position_ids, + workspace.past_seen, + num_rows, + BLOCK=row_block, + ) + except Exception as exc: + raise DSparkCompactLayoutError( + "DSpark compact-layout position snapshot failed", past_seen_valid=False + ) from exc + try: + _materialize_compact_tokens_kernel[(triton.cdiv(num_tokens, token_block),)]( + result.req_idx, + result.qo_indptr, + workspace.past_seen, + batch_slots, + new_tokens, + new_tokens_lens, + next_draft_tokens, + input_ids, + position_ids, + previous_pos_indices, + previous_pos_id_offsets, + draft_tokens, + new_tokens.stride(0), + new_tokens.stride(1), + next_draft_tokens.stride(0), + next_draft_tokens.stride(1), + num_tokens, + num_real, + real_tokens, + BLOCK=token_block, + ) + _commit_compact_rows_kernel[(1,)]( + result.verify_lens, + result.qo_indptr, + old_verify_lens, + old_qo_indptr, + kv_lens, + previous_kv_lens_offsets, + num_rows, + BLOCK=row_block, + ) + except Exception as exc: + raise DSparkCompactLayoutError( + "DSpark fused compact-layout materialization failed", past_seen_valid=True + ) from exc + return + + # Hardware-independent oracle for exact differential coverage. + split_lens = old_verify_lens[:num_rows].clone() + split_qo = old_qo_indptr[: num_rows + 1].to(torch.long) + past_seen = position_ids[split_qo[:-1]].clone() + workspace.past_seen[:num_rows].copy_(past_seen) + past_seen = workspace.past_seen[:num_rows] + req_idx = result.req_idx + flat = torch.arange(num_tokens, device=device) + offset = flat - result.qo_indptr.to(torch.long)[req_idx] + prev_slots = batch_slots[:num_real].to(torch.long) + slots_tok = prev_slots[req_idx.clamp(max=num_real - 1)] + input_ids[:num_tokens].copy_( + new_tokens.transpose(0, 1)[slots_tok, offset].flatten().to(input_ids.dtype) + ) + position_ids[:num_tokens].copy_(past_seen[req_idx] + offset.to(past_seen.dtype)) + previous_pos_indices[:real_tokens].copy_(slots_tok[:real_tokens].to(previous_pos_indices.dtype)) + previous_pos_id_offsets[:real_tokens].copy_( + new_tokens_lens[slots_tok[:real_tokens]].to(previous_pos_id_offsets.dtype) + ) + if real_draft > 0: + draft_tokens[:real_draft].copy_( + gather_packed_draft_tokens( + next_draft_tokens=next_draft_tokens, + batch_slots=prev_slots, + verify_lens=result.verify_lens, + qo_indptr=result.qo_indptr, + num_real=num_real, + total_draft_tokens=real_draft, + ).to(draft_tokens.dtype) + ) + window_delta = result.verify_lens - split_lens + kv_lens[:num_rows] += 2 * window_delta + previous_kv_lens_offsets[:num_rows] -= window_delta.to(previous_kv_lens_offsets.dtype) + old_verify_lens[:num_rows].copy_(result.verify_lens) + old_qo_indptr[: num_rows + 1].copy_(result.qo_indptr) + + def select_windows_device( *, confidence_logits: torch.Tensor, slot_idx: torch.Tensor, - num_real: torch.Tensor, - budget: torch.Tensor, + num_real: Union[int, torch.Tensor], + budget: Union[int, 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, + use_fused_exact: bool = False, + workspace: Optional[DeviceWindowWorkspace] = None, ) -> DeviceWindowResult: """Rank the batch by fresh confidence and pack it into the agreed bucket. @@ -138,9 +541,12 @@ def select_windows_device( 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. + num_real: Python integer in the production pre-replay prologue, or a + 0-d integer tensor for a device-owned caller; rows past it are + padding. + budget: Python integer in the production pre-replay prologue, or a + 0-d integer tensor paired with tensor ``num_real``; verify tokens + above the floor -- the host-decided (lagged) capacity knob. graph_num_tokens: the captured token bucket (capture constant). cfg: scheduling bounds; ``resolved_max_verify_len + 1`` is the per-row token ceiling. @@ -154,10 +560,19 @@ def select_windows_device( 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``. + use_fused_exact: enable the policy-neutral Triton fast path. + Production sets this only for shapes compiled successfully during + warmup and when its independent default-off switch is enabled; + false uses the established tensor schedule/fill implementation. + workspace: reusable output storage for the fused exact path. It is + ignored by the tensor fallback so that path remains an independent + allocation-owning oracle. Returns: - :class:`DeviceWindowResult`; every tensor is derived without a host - sync, so the call is capture-safe end to end. + :class:`DeviceWindowResult`; every output tensor stays device-resident. + The Python-control exact path is intended for the pre-replay prologue, + while the tensor-control fallback may be captured by a compatible + caller. """ if confidence_logits.dim() != 2: raise ValueError( @@ -165,6 +580,20 @@ def select_windows_device( ) device = confidence_logits.device padded_bs = slot_idx.numel() + controls_are_tensors = isinstance(num_real, torch.Tensor) + if controls_are_tensors != isinstance(budget, torch.Tensor): + raise TypeError("num_real and budget must both be Python ints or both be tensors") + if controls_are_tensors: + if num_real.dim() != 0 or budget.dim() != 0: + raise ValueError("tensor num_real and budget must both be 0-d") + # Narrow integer controls can overflow while multiplying by K or while + # casting the verification budget; production scalar controls are ordinary + # Python ints, and a device-owned caller must use a safe width. + integer_dtypes = {torch.int32, torch.int64} + if num_real.dtype not in integer_dtypes or budget.dtype not in integer_dtypes: + raise TypeError("tensor num_real and budget must both have integer dtype") + if num_real.device != device or budget.device != device: + raise ValueError("tensor num_real and budget must share confidence_logits.device") selected = confidence_logits.index_select(0, slot_idx.to(torch.long)) if stamp is not None: @@ -179,24 +608,58 @@ def select_windows_device( 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, + fused_exact = ( + use_fused_exact + and pad_len is not None + and not controls_are_tensors + and padded_bs <= 256 + and cfg.survival_eps > 0.0 ) + if fused_exact: + if workspace is not None: + workspace.validate(device=device, num_rows=padded_bs, num_tokens=graph_num_tokens) + filled = schedule_verify_lens_topk_fused_fill( + survival=survival, + budget=int(budget), + num_real=int(num_real), + pad_len=int(pad_len), + cfg=cfg, + graph_num_tokens=graph_num_tokens, + out=None if workspace is None else workspace.verify_lens, + ) + else: + # Pad rows must win no policy budget. Zero is below the configured + # numerical epsilon, preserving the established tensor scheduler + # exactly. The fused scheduler masks them internally and avoids this + # arange/where temporary in production. + 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 + 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, + ) + if fused_exact and workspace is not None: + _build_row_maps_workspace( + workspace=workspace, + num_rows=padded_bs, + graph_num_tokens=graph_num_tokens, + max_token_len=max_token_len, + ) + return DeviceWindowResult( + verify_lens=workspace.verify_lens[:padded_bs], + qo_indptr=workspace.qo_indptr[: padded_bs + 1], + req_idx=workspace.req_idx[:graph_num_tokens], + kv_correction=workspace.kv_correction[:graph_num_tokens], + workspace=workspace, + ) req_idx, correction = build_row_maps_device(filled, graph_num_tokens=graph_num_tokens) return DeviceWindowResult( verify_lens=filled, diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py index 69d5b707592d..f955fec01dca 100644 --- a/tensorrt_llm/_torch/speculative/dspark_verify.py +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -124,6 +124,8 @@ def __init__( self._prev_buffer: torch.Tensor | None = None self._prev_event: torch.cuda.Event | None = None self._prev_valid = False + self._host_confidence_stamp: torch.Tensor | None = None + self._prev_confidence_stamp: torch.Tensor | None = None # Count every path that declines to trim; degradation here is silent. self.stats = { "fallback_no_snapshot": 0, @@ -304,7 +306,11 @@ def allocate_exact_sps_candidate( def max_verify_len(self) -> int: return self.cfg.resolved_max_verify_len - def stage_confidence(self, confidence_logits: torch.Tensor) -> None: + def stage_confidence( + self, + confidence_logits: torch.Tensor, + confidence_stamp: torch.Tensor | None = None, + ) -> None: """Start a non-blocking device->host copy of the confidence buffer. ``confidence_logits`` is the worker's whole slot-indexed buffer, staged @@ -317,12 +323,24 @@ def stage_confidence(self, confidence_logits: torch.Tensor) -> None: """ if confidence_logits is None or confidence_logits.shape[0] == 0: return + # Validate before rotating any buffer; a rejected stamp must leave the + # original decode snapshot/event ownership unchanged. + if confidence_stamp is not None and ( + confidence_stamp.dtype != torch.int32 + or confidence_stamp.ndim != 1 + or confidence_stamp.shape[0] != confidence_logits.shape[0] + ): + raise ValueError("Confidence stamps must be int32, one per slot") if self.device_windows: # Rotate before overwriting: what was current becomes the older # (lag-2) snapshot the budget argmax reads; the old older buffer's # storage is reused for this staging. Pure reference swaps. self._host_buffer, self._prev_buffer = (self._prev_buffer, self._host_buffer) self._copy_event, self._prev_event = (self._prev_event, self._copy_event) + self._host_confidence_stamp, self._prev_confidence_stamp = ( + self._prev_confidence_stamp, + self._host_confidence_stamp, + ) self._prev_valid = self._snapshot_valid self._snapshot_valid = False if self._host_buffer is None or self._host_buffer.shape != confidence_logits.shape: @@ -334,6 +352,20 @@ def stage_confidence(self, confidence_logits: torch.Tensor) -> None: ) self._copy_event = torch.cuda.Event() self._host_buffer.copy_(confidence_logits, non_blocking=True) + if confidence_stamp is not None: + if ( + self._host_confidence_stamp is None + or self._host_confidence_stamp.shape != confidence_stamp.shape + ): + self._host_confidence_stamp = torch.empty( + confidence_stamp.shape, + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), + ) + self._host_confidence_stamp.copy_(confidence_stamp, non_blocking=True) + else: + self._host_confidence_stamp = None self._copy_event.record() self._snapshot_valid = True diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index e8bbc442b234..0ab7de441515 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3583,6 +3583,32 @@ def spec_dec_mode(self): return TorchSpeculativeDecodingMode.DFLASH +class DSparkMixedEvidenceConfig(StrictBaseModel): + """Optional measured BREAKABLE mixed evidence, independent of decode costs. + + All SHA256 pins come from the authenticated deployment package. The live + fingerprint must be independently observed for the current runtime; copying + the cost table's identity does not qualify it. Uncovered or invalid mixed + evidence keeps full-width native verification without disabling admitted + decode confidence scheduling. + """ + + costs_path: str = Field( + description="Measured mixed geometry/action cost JSON.") + costs_sha256: str = Field(pattern=r"^[0-9a-f]{64}$", + description="Cost-file SHA256.") + admission_path: str = Field( + description="Qualified mixed admission receipt JSON.") + admission_sha256: str = Field(pattern=r"^[0-9a-f]{64}$", + description="Receipt SHA256.") + protocol_sha256: str = Field(pattern=r"^[0-9a-f]{64}$", + description="Measured protocol SHA256.") + live_fingerprint_path: str = Field( + description="Independently generated mixed runtime fingerprint JSON.") + live_fingerprint_sha256: str = Field(pattern=r"^[0-9a-f]{64}$", + description="Live fingerprint SHA256.") + + class DSparkDecodingConfig(DecodingBaseConfig): """Configuration for DSpark speculative decoding. @@ -3592,9 +3618,8 @@ class DSparkDecodingConfig(DecodingBaseConfig): backbone forward, but it additionally refines the per-position draft logits with a lightweight sequential head (a low-rank Markov head, optionally an RNN head). A confidence head can score each draft position for verification - scheduling without truncating the proposed block. This configuration selects - native full-block drafting and does not expose confidence-scheduling or - calibration controls. + scheduling without truncating the proposed block. Confidence scheduling is + opt-in; the default configuration retains native full-block verification. Key features: - Target-dependent: captures hidden states from ``target_layer_ids``. @@ -3636,6 +3661,89 @@ class DSparkDecodingConfig(DecodingBaseConfig): "read from the draft model config (dspark_markov_head_type), " "defaulting to \"vanilla\".") + enable_confidence_scheduling: bool = Field( + default=False, + description= + "Score each drafted position with the draft's confidence head and let " + "the verification scheduler size the batch's verify budget from those " + "scores. The block is always drafted in full; only the number of tokens " + "sent to the target for verification changes, so acceptance (and hence " + "output distribution) is unaffected. Requires a checkpoint with trained " + "confidence-head weights. Gains only appear at high concurrency -- at " + "small batch sizes the step cost is nearly flat in the verified token " + "count, so trimming saves nothing and costs acceptance length.", + status="prototype") + + validate_confidence_artifact_freshness: bool = Field( + default=False, + description= + "Require exact live cost-artifact identity and a positive workload " + "admission receipt before enabling confidence scheduling. By default, " + "explicitly enabled scheduling reuses structurally valid cost and " + "calibration artifacts, with a warning that performance predictions " + "may be stale. This never changes target verification or acceptance.", + status="prototype") + + confidence_sts_path: Optional[str] = Field( + default=None, + description= + "Path to a JSON file with per-position sequential-temperature-scaling " + "(STS) values used to calibrate the confidence head. The scheduler " + "consumes a cumulative product of per-position probabilities, so a " + "per-position calibration error compounds geometrically along the " + "block. If None, the raw sigmoid is used (identity calibration).", + status="prototype") + + confidence_sps_table_path: Optional[str] = Field( + default=None, + description= + "Path to an authenticated schema-v2 JSON table of measured decode step " + "costs keyed by exact rank-local CUDA-graph shape (G, V).", + status="prototype") + + confidence_sps_live_fingerprint_path: Optional[str] = Field( + default=None, + description= + "Path to an independently generated JSON fingerprint of the active " + "runtime. Required when validate_confidence_artifact_freshness=True. " + "Otherwise it is optional observed provenance, kept separate from " + "the cost artifact's original metadata.", + status="prototype") + + confidence_admission_receipt_path: Optional[str] = Field( + default=None, + description= + "Path to an authenticated workload-level confidence admission receipt. " + "Required only when validate_confidence_artifact_freshness=True; it " + "must prove positive aggregate compact value after the fixed " + "confidence-path overhead and safety margin are charged.", + status="prototype") + + confidence_admission_receipt_sha256: Optional[str] = Field( + default=None, + description="Externally authenticated lowercase SHA256 of " + "confidence_admission_receipt_path. This pin must come from the sealed " + "deployment package rather than from the receipt itself.", + status="prototype") + + enable_fused_confidence_scheduler: bool = Field( + default=False, + description= + "Use the optional policy-neutral Triton implementation of confidence " + "top-k allocation plus exact CUDA-graph bucket fill. This changes only " + "the scheduler launch path: budgets, survival thresholds, selected " + "windows, and fallback behavior remain identical. It is independently " + "default-off and requires confidence scheduling.", + status="prototype") + + confidence_mixed_evidence: Optional[DSparkMixedEvidenceConfig] = Field( + default=None, + description="Optional separately admitted physical-K3, eight-rank " + "BREAKABLE mixed scheduling evidence. Does not reuse decode T(G,V) costs. " + "None, unsupported routes, stale ownership, or uncovered mixed geometry " + "retain native mixed verification while admitted decode remains enabled.", + status="prototype") + decoding_type: Literal["DSpark"] = Field(default="DSpark") attention_backend: Literal["AUTO", "VANILLA", "TRTLLM", "CUTEDSL"] = Field( @@ -3651,6 +3759,51 @@ class DSparkDecodingConfig(DecodingBaseConfig): "and needs a cute-dsl MLA decode taking per-token kv_bounds that is not " "upstream yet. Which kernel each name selects: MLADSparkForCausalLM.") + @model_validator(mode="after") + def validate_confidence_scheduling(self): + if not self.enable_confidence_scheduling: + if (self.validate_confidence_artifact_freshness + or self.enable_fused_confidence_scheduler + or self.confidence_sts_path + or self.confidence_sps_table_path + or self.confidence_sps_live_fingerprint_path + or self.confidence_admission_receipt_path + or self.confidence_admission_receipt_sha256 + or self.confidence_mixed_evidence is not None): + raise ValueError( + "validate_confidence_artifact_freshness / " + "enable_fused_confidence_scheduler / confidence_sts_path / " + "confidence_sps_table_path / " + "confidence_sps_live_fingerprint_path / " + "confidence_admission_receipt_path / " + "confidence_admission_receipt_sha256 / confidence_mixed_evidence require " + "enable_confidence_scheduling=True") + return self + + if self.max_draft_len == 1: + # Physical K1 has no smaller verifier window, so confidence + # scheduling can only add control-path work. Resolve it to the + # ordinary feature-off configuration before any runtime component + # observes the config. This preserves the original static K path + # rather than running a confidence-enabled native fallback. + logger.warning( + "DSpark confidence scheduling has no compact tier below the " + "native K1 layout; using the feature-off static K1 path and " + "ignoring confidence-scheduling inputs.") + self.enable_confidence_scheduling = False + self.validate_confidence_artifact_freshness = False + self.enable_fused_confidence_scheduler = False + self.confidence_sts_path = None + self.confidence_sps_table_path = None + self.confidence_sps_live_fingerprint_path = None + self.confidence_admission_receipt_path = None + self.confidence_admission_receipt_sha256 = None + self.confidence_mixed_evidence = None + return self + # Resolve optional freshness admission before runtime components + # snapshot the flag. Exact graph/K coverage is checked by the engine. + return self + @model_validator(mode="after") def set_max_total_draft_tokens(self): self.max_total_draft_tokens = self.max_draft_len @@ -6900,6 +7053,170 @@ def validate_encode_only_torch_compile_config(self) -> 'TorchLlmArgs': "graphs or disable enable_piecewise_cuda_graph.") return self + def _validate_dspark_confidence_environment(self) -> None: + """Reject configurations confidence scheduling cannot honour. + + Every check here guards a path that would otherwise degrade *silently*: + the run comes up, produces plausible text, and either never takes the + ragged path at all or takes it with a kernel that mis-attributes rows. + A hard error at construction is the only signal that survives. + """ + if not self.speculative_config.draft_is_embedded_in_target: + raise ValueError( + "DSpark confidence scheduling supports only the embedded " + "DeepSeek-V4 draft. Standalone DSpark retains native " + "verification; disable confidence scheduling for that path.") + cuda_graph_config = self.cuda_graph_config + if self.pipeline_parallel_size > 1: + raise ValueError( + "speculative_config.enable_confidence_scheduling=True does not yet " + "support pipeline parallelism. The PP executor's pipelined " + "microbatches do not provide the single-iteration host-input " + "lifetime boundary required by ragged verification staging. " + "Set pipeline_parallel_size=1 or disable confidence scheduling." + ) + if self.guided_decoding_backend is not None: + raise ValueError( + "speculative_config.enable_confidence_scheduling=True does not yet " + "support guided decoding. Guided-decoder bitmasks are laid " + "out with one uniform speculative stride; applying them to " + "packed per-request verify windows would shift every request " + "after the first short window. Disable guided decoding or " + "disable confidence scheduling.") + if self.enable_lora or self.lora_config is not None: + raise ValueError( + "speculative_config.enable_confidence_scheduling=True does not yet " + "support LoRA. LoRA metadata expands generation requests " + "with one uniform speculative stride, which cannot represent " + "packed per-request verify windows. Disable LoRA or disable " + "confidence scheduling.") + # `fit_ragged_verify_lens` plans against the padded row count and + # reserves tokens for the CUDA-graph padding rows, so it presumes those + # rows exist. With padding off, `_get_padded_batch` adds none: the batch + # then carries a token total that keys into a graph nobody captured and + # drops to eager on every ragged step -- which costs far more than the + # trimmed tokens save -- while real requests were still shrunk to leave + # room for padding that never arrived. + if cuda_graph_config is None: + raise ValueError( + "speculative_config.enable_confidence_scheduling=True requires " + "cuda_graph_config to be set with enable_padding=True. Ragged " + "verification fits each step's token total onto a captured " + "CUDA-graph bucket; with no graphs there is no bucket to fit " + "and every step would silently run eager.") + if not cuda_graph_config.enable_padding: + raise ValueError( + "speculative_config.enable_confidence_scheduling=True requires " + "cuda_graph_config.enable_padding=True (it defaults to False). " + "The ragged token budget is planned against the padded batch " + "size and reserves tokens for the padding rows; without padding " + "those rows never appear, so the step's token total misses its " + "captured bucket and falls out of graph replay.") + + sparse_attention_config = self.sparse_attention_config + if sparse_attention_config is not None: + # Both recover "which request is this row" from a fixed next_n -- + # GVR indexes its previous-step top-k hint by request and strides it + # by next_n, and the CuTe-DSL kernel takes next_n as a compile-time + # constant and JIT cache key. Neither can express per-request + # windows. The C++ launcher already refuses the combination + # (indexerTopK.cu, "ragged rowKvLens is incompatible with the GVR + # preIdx hint"), but that fires deep in a decode step; say so here. + incompatible = [ + name for name in ("enable_heuristic_topk", "use_cute_dsl_topk") + if getattr(sparse_attention_config, name, False) + ] + if incompatible: + raise ValueError( + "speculative_config.enable_confidence_scheduling=True is " + "incompatible with sparse_attention_config." + f"{' / '.join(incompatible)}=True: those top-k paths " + "reconstruct each row's request and window position from a " + "single batch-wide next_n, which a ragged batch does not " + "have. Disable them, or disable confidence scheduling.") + # Not an error: the DSL paged-MQA-logits path is structurally + # uniform too, but the expanded (one row per query token) path it + # falls back to is numerically equivalent, and it is the default on + # DeepSeek-V4 anyway. Still worth saying out loud, since the knob + # will read as honoured otherwise. + if getattr(sparse_attention_config, "use_cute_dsl_paged_mqa_logits", + False): + logger.warning( + "sparse_attention_config.use_cute_dsl_paged_mqa_logits=True " + "is ignored on ragged-verification steps: the DSL kernel " + "broadcasts one KV length per request across a fixed next_n. " + "Those steps use the expanded DeepGEMM path instead, which " + "is numerically equivalent.") + + @staticmethod + def _resolve_dspark_confidence_workload_admission( + spec_cfg: DSparkDecodingConfig) -> None: + """Resolve confidence activation before engines or workers observe it.""" + from tensorrt_llm._torch.speculative.dspark_planner import \ + load_confidence_workload_admission + + receipt_path = spec_cfg.confidence_admission_receipt_path + receipt_sha256 = spec_cfg.confidence_admission_receipt_sha256 + sps_table_path = spec_cfg.confidence_sps_table_path + live_fingerprint_path = spec_cfg.confidence_sps_live_fingerprint_path + if not spec_cfg.validate_confidence_artifact_freshness: + if not sps_table_path: + raise ValueError( + "DSpark confidence scheduling requires confidence_sps_table_path; " + "a missing measured cost table cannot be timed as enabled.") + logger.warning( + "DSpark confidence scheduling is enabled with reusable artifacts. " + "Cost/calibration predictions and performance economics are not " + "freshness-validated and may be stale; target acceptance is unchanged. " + "Set validate_confidence_artifact_freshness=True to require exact " + "live provenance and positive workload admission.") + return + missing = [ + name for name, value in ( + ("confidence_admission_receipt_path", receipt_path), + ("confidence_admission_receipt_sha256", receipt_sha256), + ("confidence_sps_table_path", sps_table_path), + ("confidence_sps_live_fingerprint_path", live_fingerprint_path), + ) if not value + ] + fallback_reason = ("missing " + ", ".join(missing)) if missing else None + if fallback_reason is None: + assert receipt_path is not None + assert receipt_sha256 is not None + assert sps_table_path is not None + assert live_fingerprint_path is not None + try: + admission = load_confidence_workload_admission( + receipt_path, + expected_receipt_sha256=receipt_sha256, + sps_cost_table_path=sps_table_path, + live_engine_fingerprint_path=live_fingerprint_path, + physical_k=spec_cfg.max_draft_len, + ) + if not admission.admitted: + fallback_reason = ( + "the authenticated workload replay does not have positive " + "net compact value after fixed overhead and safety margin" + ) + except (OSError, TypeError, ValueError) as error: + fallback_reason = str(error) + + if fallback_reason is None: + return + + logger.warning( + f"DSpark confidence scheduling is not deployment-admitted ({fallback_reason}); " + f"using the feature-off static K{spec_cfg.max_draft_len} path.") + spec_cfg.enable_confidence_scheduling = False + spec_cfg.validate_confidence_artifact_freshness = False + spec_cfg.enable_fused_confidence_scheduler = False + spec_cfg.confidence_sts_path = None + spec_cfg.confidence_sps_table_path = None + spec_cfg.confidence_sps_live_fingerprint_path = None + spec_cfg.confidence_admission_receipt_path = None + spec_cfg.confidence_admission_receipt_sha256 = None + spec_cfg.confidence_mixed_evidence = None + @model_validator(mode="after") def normalize_prefill_cuda_graph_config(self) -> 'TorchLlmArgs': """Normalize legacy piecewise CUDA graph options into prefill fields.""" @@ -7379,6 +7696,10 @@ def _dspark_get( "DSpark block_size must equal max_draft_len; got " f"block_size={spec_cfg.block_size} and " f"max_draft_len={spec_cfg.max_draft_len}") + if spec_cfg.enable_confidence_scheduling: + self._resolve_dspark_confidence_workload_admission(spec_cfg) + if spec_cfg.enable_confidence_scheduling: + self._validate_dspark_confidence_environment() if isinstance(self.speculative_config, SADecodingConfig): pool_size = self.speculative_config.global_pool_size diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index b96235b86327..cd6bbd97f36b 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1679,6 +1679,16 @@ "kind": "categorical", "path": "speculative_config.eagle3_model_arch" }, + { + "capture_policy": "bool", + "kind": "value", + "path": "speculative_config.enable_confidence_scheduling" + }, + { + "capture_policy": "bool", + "kind": "value", + "path": "speculative_config.enable_fused_confidence_scheduler" + }, { "capture_policy": "bool", "kind": "value", @@ -1838,6 +1848,11 @@ "kind": "value", "path": "speculative_config.use_relaxed_acceptance_for_thinking" }, + { + "capture_policy": "bool", + "kind": "value", + "path": "speculative_config.validate_confidence_artifact_freshness" + }, { "capture_policy": "int", "kind": "value", diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index c77dff8414ac..54fa5dd947e2 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -193,6 +193,7 @@ l0_cpu: - unittest/llmapi/test_whisper_suppress_tokens_processor.py - unittest/llmapi/test_kv_cache_dtype_override.py - unittest/llmapi/test_llm_args.py + - unittest/llmapi/test_dspark_workload_admission_config.py - unittest/llmapi/test_mx_args.py - unittest/llmapi/test_custom_tokenizer_aliases.py - unittest/llmapi/test_llm_quant.py diff --git a/tests/unittest/_torch/executor/engine/test_metadata.py b/tests/unittest/_torch/executor/engine/test_metadata.py index 801d01b70f2f..392afda5f521 100644 --- a/tests/unittest/_torch/executor/engine/test_metadata.py +++ b/tests/unittest/_torch/executor/engine/test_metadata.py @@ -95,3 +95,47 @@ def test_build_attention_metadata_forwards_shared_and_cache_inputs() -> None: assert metadata.draft_kv_cache_manager is draft_kv_cache_manager assert metadata.enable_context_mla_with_cached_kv assert metadata.num_heads_per_kv == 4 + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_build_attention_metadata_passes_ragged_opt_in_to_supported_backend(enabled) -> None: + class RaggedMetadata(_AttentionMetadata): + enable_ragged_verification = False + + model_config = SimpleNamespace( + pretrained_config=SimpleNamespace(), + sparse_attention_config=None, + enable_flash_mla=False, + ) + metadata = build_attention_metadata( + model_config, + max_batch_size=4, + max_num_tokens=16, + max_beam_width=1, + attention_backend=SimpleNamespace(Metadata=RaggedMetadata), + attention_runtime_features=AttentionRuntimeFeatures(), + mapping=object(), + cache_indirection=None, + enable_ragged_verification=enabled, + ) + assert metadata.enable_ragged_verification is enabled + + +def test_build_attention_metadata_does_not_pass_ragged_option_to_other_backend() -> None: + model_config = SimpleNamespace( + pretrained_config=SimpleNamespace(), + sparse_attention_config=None, + enable_flash_mla=False, + ) + metadata = build_attention_metadata( + model_config, + max_batch_size=4, + max_num_tokens=16, + max_beam_width=1, + attention_backend=_AttentionBackend, + attention_runtime_features=AttentionRuntimeFeatures(), + mapping=object(), + cache_indirection=None, + enable_ragged_verification=True, + ) + assert not hasattr(metadata, "enable_ragged_verification") diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py index c6509d5c503e..9cbdbaddfe50 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py @@ -8,6 +8,7 @@ import tensorrt_llm import tensorrt_llm.bindings +from tensorrt_llm._torch.pyexecutor._util import _get_num_reserved_index_slots from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2 as kv_cache_v2_module from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState, SamplingConfig @@ -16,6 +17,27 @@ CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType +@pytest.mark.parametrize( + ("spec_config", "estimating_kv_cache", "expected"), + [ + (None, False, 1), + (SimpleNamespace(enable_confidence_scheduling=False), False, 1), + (SimpleNamespace(enable_confidence_scheduling=True), False, 2), + (SimpleNamespace(enable_confidence_scheduling=True), True, 0), + ], + ids=[ + "no-speculation", + "static-dspark", + "confidence-dspark", + "confidence-estimator", + ], +) +def test_index_mapper_reservation_tracks_confidence_scheduling( + spec_config: SimpleNamespace | None, estimating_kv_cache: bool, expected: int +) -> None: + assert _get_num_reserved_index_slots(spec_config, estimating_kv_cache) == expected + + def _manager( *, is_draft: bool, 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 index 82ca8377d59d..2846cb2f50dc 100644 --- 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 @@ -14,6 +14,8 @@ from tensorrt_llm._torch.speculative import dspark as dspark_module from tensorrt_llm._torch.speculative.dspark import DSparkSpecMetadata, DSv4DSparkWorker +pytestmark = pytest.mark.cpu_only + @pytest.mark.parametrize("mode", ["Eagle", "MTP"]) def test_tree_speculation_keeps_native_dummy_identity(mode): @@ -64,6 +66,8 @@ def allocate(*args, **kwargs): worker._attention_warmup_attempted = True worker.return_confidence = False worker._confidence_logits = None + worker._mixed_confidence_tracking = False + worker._confidence_incarnations = {} draft = SimpleNamespace( block_size=physical_k, num_stages=1, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_admission_resolution.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_admission_resolution.py new file mode 100644 index 000000000000..e35300747444 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_admission_resolution.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Configuration-level native fallback, before any worker snapshots its flags.""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.speculative import dspark_planner +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig, TorchLlmArgs + +pytestmark = pytest.mark.cpu_only + +_ARTIFACT_FIELDS = ( + "confidence_sts_path", + "confidence_sps_table_path", + "confidence_sps_live_fingerprint_path", + "confidence_admission_receipt_path", + "confidence_admission_receipt_sha256", +) + + +def _config(physical_k=3, **kwargs): + return DSparkDecodingConfig( + max_draft_len=physical_k, + enable_confidence_scheduling=True, + validate_confidence_artifact_freshness=True, + enable_fused_confidence_scheduler=True, + **kwargs, + ) + + +@pytest.mark.parametrize("outcome", ["missing", "negative", "stale", "unreadable", "positive"]) +@pytest.mark.parametrize("physical_k", [2, 3, 4, 5, 8]) +def test_admission_resolves_before_worker_construction(monkeypatch, physical_k, outcome): + artifacts = {name: "/sealed/artifact.json" for name in _ARTIFACT_FIELDS} + artifacts["confidence_admission_receipt_sha256"] = "a" * 64 + config = _config(physical_k, **({} if outcome == "missing" else artifacts)) + calls = [] + + def load(*args, **kwargs): + calls.append((args, kwargs)) + if outcome == "stale": + raise ValueError("source identity mismatch") + if outcome == "unreadable": + raise OSError("artifact unavailable") + return SimpleNamespace(admitted=outcome == "positive") + + monkeypatch.setattr(dspark_planner, "load_confidence_workload_admission", load) + TorchLlmArgs._resolve_dspark_confidence_workload_admission(config) + assert config.max_draft_len == physical_k + assert config.tokens_per_gen_step == physical_k + 1 + assert config.enable_confidence_scheduling is (outcome == "positive") + assert config.enable_fused_confidence_scheduler is (outcome == "positive") + # Experiment classification must use the resolved snapshot, not the input + # enable flag. A successful fallback is not an enabled-native experiment. + resolved = config.model_dump() + assert resolved["enable_confidence_scheduling"] is (outcome == "positive") + assert resolved["enable_fused_confidence_scheduler"] is (outcome == "positive") + for name in _ARTIFACT_FIELDS: + assert getattr(config, name) == (artifacts[name] if outcome == "positive" else None) + assert resolved[name] == (artifacts[name] if outcome == "positive" else None) + if outcome == "missing": + assert not calls + else: + assert calls[0][1]["physical_k"] == physical_k + assert calls[0][1]["expected_receipt_sha256"] == "a" * 64 + + +def test_k1_never_needs_admission_or_confidence_artifacts(): + config = _config(1) + assert config.enable_confidence_scheduling is False + assert config.enable_fused_confidence_scheduler is False + assert config.tokens_per_gen_step == 2 + assert all(getattr(config, name) is None for name in _ARTIFACT_FIELDS) + + +def test_explicit_enable_defaults_to_reuse_without_inventing_admission(monkeypatch): + config = DSparkDecodingConfig( + max_draft_len=3, + enable_confidence_scheduling=True, + confidence_sps_table_path="/original/sps.json", + confidence_sts_path="/original/sts.json", + ) + + def reject_load(*args, **kwargs): + raise AssertionError("Reuse must not fabricate or require an admission receipt") + + monkeypatch.setattr(dspark_planner, "load_confidence_workload_admission", reject_load) + TorchLlmArgs._resolve_dspark_confidence_workload_admission(config) + assert config.enable_confidence_scheduling is True + assert config.validate_confidence_artifact_freshness is False + assert config.confidence_admission_receipt_path is None + assert config.confidence_sps_table_path == "/original/sps.json" + restored = DSparkDecodingConfig.model_validate_json(config.model_dump_json()) + assert restored.enable_confidence_scheduling is True + assert restored.validate_confidence_artifact_freshness is False + assert restored.confidence_sts_path == "/original/sts.json" + + +def test_reuse_refuses_to_claim_enabled_without_a_cost_table(): + config = DSparkDecodingConfig(max_draft_len=3, enable_confidence_scheduling=True) + with pytest.raises(ValueError, match="requires confidence_sps_table_path"): + TorchLlmArgs._resolve_dspark_confidence_workload_admission(config) + + +def test_default_off_and_strict_mode_round_trip(): + config = DSparkDecodingConfig(max_draft_len=3) + assert config.enable_confidence_scheduling is False + assert config.validate_confidence_artifact_freshness is False + strict = _config(3) + restored = DSparkDecodingConfig.model_validate_json(strict.model_dump_json()) + assert restored.validate_confidence_artifact_freshness is True + with pytest.raises(ValueError, match="enable_confidence_scheduling=True"): + DSparkDecodingConfig(max_draft_len=3, validate_confidence_artifact_freshness=True) + + +@pytest.mark.parametrize("field_name", _ARTIFACT_FIELDS) +def test_feature_off_rejects_confidence_only_inputs(field_name): + with pytest.raises(ValueError, match="enable_confidence_scheduling=True"): + DSparkDecodingConfig(max_draft_len=3, **{field_name: "unconsumed"}) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_artifact_reuse.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_artifact_reuse.py new file mode 100644 index 000000000000..65d94ac666d0 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_artifact_reuse.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Cost reuse changes predictions/activation, never target acceptance.""" + +import copy +import hashlib +import json +from pathlib import Path + +import pytest +from pydantic import BaseModel + +from tensorrt_llm._torch.speculative.dspark_planner import ( + ExactSpsCostRow, + ExactSpsCostTable, + ExactSpsDrainGuard, + load_runtime_sps_cost_table, + select_exact_sps_candidate, +) +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig +from tensorrt_llm.usage.llmapi_config import collect_llm_api_config_payloads, manifest_rows + +pytestmark = pytest.mark.cpu_only + + +def _payload(): + fingerprint = { + "gpu": "B300", + "gpu_count": 8, + "gpu_snapshot_sha256": "a" * 64, + "global_graph_batch_sizes": [512], + "max_draft_len": 3, + "rank_local_graph_batch_sizes": [64], + "runtime_snapshot": "original-runtime", + "source_diff_sha256": "b" * 64, + "source_head": "original-source", + "topology": "DEP8", + } + return { + "schema_version": 2, + "minimum_predicted_gain": 0.01, + "engine_fingerprint": fingerprint, + "engine_fingerprint_sha256": hashlib.sha256( + json.dumps(fingerprint, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest(), + "cost_tables": {"64": {"token_counts": [0, 128], "step_time_ms": [10.0, 8.0]}}, + "measurements": [ + { + "rank_local_graph_batch_size": 64, + "rank_local_verifier_budget": budget, + "step_time_ms": step, + "source_result_sha256": "c" * 64, + } + for budget, step in ((0, 10.0), (128, 8.0)) + ], + } + + +def _load(tmp_path, payload, *, strict=False, live=None): + path = tmp_path / "sps.json" + path.write_text(json.dumps(payload), encoding="utf-8") + live_path = None + if live is not None: + live_path = tmp_path / "live.json" + live_path.write_text(json.dumps(live), encoding="utf-8") + return load_runtime_sps_cost_table( + path, + graph_batch_sizes=[64], + max_draft_len=3, + live_engine_fingerprint_path=live_path, + validate_freshness=strict, + ) + + +def _select(table, *, yield_value=9.0): + return select_exact_sps_candidate( + graph_batch_size=64, + native_expected_yield=10.0, + compact_expected_yields={128: yield_value}, + compact_max_yield_losses_per_request={128: 10.0 - yield_value}, + cost_table=table, + ) + + +def test_reuse_without_live_receipt_or_drain_can_genuinely_select_compact(tmp_path): + payload = _payload() + before = copy.deepcopy(payload) + table, metadata = _load(tmp_path, payload) + assert metadata == before + assert table.validate_artifact_freshness is False + assert table.iteration_drain_guard is None + assert _select(table) == 128 + assert _select(table, yield_value=8.0) == 0 + assert table.minimum_predicted_gain == 0.01 + + +@pytest.mark.parametrize( + "field", ["source_head", "runtime_snapshot", "source_diff_sha256", "gpu_snapshot_sha256"] +) +def test_reuse_retains_original_metadata_but_strict_rejects_staleness(tmp_path, field): + payload = _payload() + live = copy.deepcopy(payload["engine_fingerprint"]) + live[field] = "d" * 64 + table, metadata = _load(tmp_path, payload, live=live) + assert _select(table) == 128 + assert metadata["engine_fingerprint"] == payload["engine_fingerprint"] + assert metadata["engine_fingerprint"][field] != live[field] + with pytest.raises(ValueError, match="does not match active runtime"): + _load(tmp_path, payload, strict=True, live=live) + + +def test_strict_requires_independent_live_and_retains_absent_drain_refusal(tmp_path): + payload = _payload() + with pytest.raises(ValueError, match="independently generated"): + _load(tmp_path, payload, strict=True) + table, _ = _load( + tmp_path, payload, strict=True, live=copy.deepcopy(payload["engine_fingerprint"]) + ) + assert _select(table) == 0 + + +@pytest.mark.parametrize( + "fault", + [ + "nonfinite", + "missing_native", + "illegal_window", + "missing_graph", + "wrong_k", + "bad_hash", + "missing_provenance", + ], +) +def test_reuse_keeps_all_structural_guards(tmp_path, fault): + payload = _payload() + if fault == "nonfinite": + payload["cost_tables"]["64"]["step_time_ms"][1] = float("nan") + if fault == "missing_native": + payload["cost_tables"]["64"]["token_counts"][0] = 64 + if fault == "illegal_window": + payload["cost_tables"]["64"]["token_counts"][1] = 257 + if fault == "missing_graph": + payload["cost_tables"] = {"32": payload["cost_tables"]["64"]} + if fault == "wrong_k": + payload["engine_fingerprint"]["max_draft_len"] = 5 + payload["engine_fingerprint_sha256"] = hashlib.sha256( + json.dumps( + payload["engine_fingerprint"], sort_keys=True, separators=(",", ":") + ).encode() + ).hexdigest() + if fault == "bad_hash": + payload["engine_fingerprint_sha256"] = "0" * 64 + if fault == "missing_provenance": + payload["measurements"].pop() + with pytest.raises((TypeError, ValueError)): + _load(tmp_path, payload) + + +def test_provided_drain_policy_prices_and_margin_are_unchanged_in_both_modes(): + kwargs = { + "tables": {64: ExactSpsCostRow((0, 128), (10.0, 8.0))}, + "max_draft_len": 3, + "minimum_predicted_gain": 0.01, + "iteration_drain_guard": ExactSpsDrainGuard(1.0, 1.0, 0.0, 64, "a" * 64), + } + strict = ExactSpsCostTable(**kwargs) + reuse = ExactSpsCostTable(**kwargs, validate_artifact_freshness=False) + assert _select(strict) == _select(reuse) == 0 + assert strict.iteration_drain_guard == reuse.iteration_drain_guard + assert strict.tables == reuse.tables + assert strict.minimum_predicted_gain == reuse.minimum_predicted_gain + assert strict.collective_identity_words != reuse.collective_identity_words + assert len(strict.collective_identity_words) == len(reuse.collective_identity_words) == 8 + + +@pytest.mark.parametrize("strict", [False, True]) +def test_freshness_mode_serialization_and_telemetry_capture_only_the_bool(strict): + class Args(BaseModel): + speculative_config: DSparkDecodingConfig + + config = DSparkDecodingConfig( + max_draft_len=3, + enable_confidence_scheduling=True, + validate_confidence_artifact_freshness=strict, + confidence_sps_table_path="/private/original/sps.json", + confidence_sts_path="/private/original/sts.json", + ) + args = Args.model_validate_json(Args(speculative_config=config).model_dump_json()) + field_path = "speculative_config.validate_confidence_artifact_freshness" + expected = {"path": field_path, "kind": "value", "capture_policy": "bool"} + assert expected in manifest_rows(Args) + captured, metadata = collect_llm_api_config_payloads(args) + assert json.loads(metadata)["capture_succeeded"] is True + assert json.loads(captured)[field_path] is strict + assert "private" not in captured + golden_path = Path(__file__).parents[5] / "tensorrt_llm/usage/llm_args_golden_manifest.json" + assert expected in json.loads(golden_path.read_text(encoding="utf-8"))["TorchLlmArgs"] 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 0416bb58956b..6b8a670def45 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 @@ -4,6 +4,7 @@ from types import SimpleNamespace +import numpy as np import pytest import torch @@ -301,3 +302,84 @@ def test_cached_shape_agreement_rejects_a_changed_ragged_key(): with pytest.raises(RuntimeError, match="no longer matches"): runner._agree_graph_shape(batch, True, frozenset()) + + +def _tier_lookup_runner(batch, *, confidence_enabled, in_graph_sampling, peer_tier): + runner = _runner_stub() + payloads = [] + + def gather(payload): + payloads.append(tuple(payload)) + peer = list(payload) + peer[-1] = peer_tier.value + return np.asarray([payload, peer], dtype=np.int64) + + runner.config = SimpleNamespace( + enable_attention_dp=True, + mapping=SimpleNamespace(tp_size=2), + dist=SimpleNamespace(tp_allgather_int64=gather), + use_mrope=False, + enable_in_graph_sampling=in_graph_sampling, + ) + runner.enabled = True + runner.enable_in_graph_sampling = in_graph_sampling + runner.is_encoder_decoder = False + runner.enable_encoder_decoder_mixed_cuda_graph = False + runner._dspark_confidence_enabled = confidence_enabled + runner._dspark_trims_submitted_tokens = confidence_enabled + runner.adp_shape_agreement = None + runner.agreed_ragged_bucket = 16 if confidence_enabled else None + runner.sparse_config = None + runner._capture_sample_type = None + runner._sample_type_resolver = lambda _batch, _override: SampleType.FAST + runner._capture_allowed = False + runner.graph_metadata = { + KeyType( + batch_size=4, + draft_len=5, + is_first_draft=False, + sample_type=tier, + ragged_verify_bucket=runner.agreed_ragged_bucket, + ): {"attn_metadata": tier, "spec_metadata": None} + for tier in (SampleType.FULL, SampleType.FAST) + } + return runner, payloads + + +@pytest.mark.parametrize("confidence_enabled", [False, True]) +@pytest.mark.parametrize("peer_tier", [SampleType.FAST, SampleType.FULL]) +def test_graph_lookup_agrees_sampling_tier_without_aliasing_ragged_bucket( + confidence_enabled, peer_tier +): + batch = _Batch([SimpleNamespace(py_draft_tokens=[1] * 5, py_verify_len=3) for _ in range(4)]) + runner, payloads = _tier_lookup_runner( + batch, confidence_enabled=confidence_enabled, in_graph_sampling=True, peer_tier=peer_tier + ) + + metadata, _, key = runner.maybe_get_cuda_graph(batch, True, SimpleNamespace()) + + assert metadata is peer_tier + assert key.sample_type is peer_tier + assert payloads == [ + (True, 4, 16, 5, SampleType.FAST.value) + if confidence_enabled + else (True, 4, SampleType.FAST.value) + ] + + +@pytest.mark.parametrize("in_graph_sampling", [False, True]) +def test_cached_shape_does_not_replace_in_graph_sampling_tier_agreement(in_graph_sampling): + batch = _Batch([SimpleNamespace(py_draft_tokens=[1] * 5, py_verify_len=3) for _ in range(4)]) + runner, payloads = _tier_lookup_runner( + batch, + confidence_enabled=True, + in_graph_sampling=in_graph_sampling, + peer_tier=SampleType.FAST, + ) + runner.adp_shape_agreement = _agreement(batch, peer_sizes=(4, 4)) + + metadata, _, key = runner.maybe_get_cuda_graph(batch, True, SimpleNamespace()) + + assert metadata is (SampleType.FAST if in_graph_sampling else SampleType.FULL) + assert key.sample_type is metadata + assert len(payloads) == int(in_graph_sampling) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_dummy_slots_cpu.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_dummy_slots_cpu.py new file mode 100644 index 000000000000..7b5fecb3d949 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_dummy_slots_cpu.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Host slot-ownership regression with CPU tensors in place of CUDA allocations.""" + +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm._torch.speculative.dspark as dspark +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import cuda_graph_dummy_request_id +from tensorrt_llm._torch.speculative.dspark import DSparkSpecMetadata, DSv4DSparkWorker + +pytestmark = pytest.mark.cpu_only + + +@pytest.mark.parametrize("physical_k", range(1, 9)) +def test_all_padding_families_preserve_real_slot_ownership(monkeypatch, physical_k): + # Exercise the real initialization and host prepare methods. This deliberately + # does not validate device copies, graph replay, or target KV-cache contents. + original_zeros = torch.zeros + original_empty = torch.empty + original_full = torch.full + + def cpu_zeros(*args, **kwargs): + kwargs["device"] = "cpu" + return original_zeros(*args, **kwargs) + + def cpu_empty(*args, **kwargs): + kwargs["device"] = "cpu" + return original_empty(*args, **kwargs) + + def cpu_full(*args, **kwargs): + kwargs["device"] = "cpu" + return original_full(*args, **kwargs) + + monkeypatch.setattr(torch, "zeros", cpu_zeros) + monkeypatch.setattr(torch, "empty", cpu_empty) + monkeypatch.setattr(torch, "full", cpu_full) + monkeypatch.setattr(dspark, "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 + worker._mixed_confidence_tracking = False + worker._confidence_incarnations = {} + draft = SimpleNamespace( + block_size=physical_k, + num_stages=1, + _attn_params={"window_size": 4, "head_dim": 2}, + ) + DSv4DSparkWorker._lazy_init(worker, draft, SimpleNamespace(max_num_requests=4)) + meta = 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 = cuda_graph_dummy_request_id( + runtime_k, variant=variant, max_draft_len=physical_k + ) + # Include ADP idle, the secondary zero-real family, and a live row. + meta.request_ids = [1000, 0, dummy] + meta.num_generations = 3 + DSparkSpecMetadata.prepare(meta) + 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] + # Removing a real request releases exactly its slot, never a dummy's. + meta.request_ids = [dummy, 0] + meta.num_generations = 2 + DSparkSpecMetadata.prepare(meta) + 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_embedded_confidence_loading.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_embedded_confidence_loading.py new file mode 100644 index 000000000000..01881356f6e9 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_embedded_confidence_loading.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Embedded-wrapper validation through the real shared parameter walk.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch.models import modeling_dspark, modeling_dspark_v41 +from tensorrt_llm._torch.models.modeling_speculative import DSparkConfidenceHead + +pytestmark = pytest.mark.cpu_only + + +def _wrapper(monkeypatch, v41): + module = modeling_dspark_v41 if v41 else modeling_dspark + cls = module.DSv41DSparkForCausalLM if v41 else module.DSv4DSparkForCausalLM + wrapper = object.__new__(cls) + nn.Module.__init__(wrapper) + stage = nn.Module() + stage.main_proj = nn.Linear(2, 2, bias=False) + stage.main_norm = nn.RMSNorm(2) + stage.norm = nn.RMSNorm(2) + stage.markov_head = None + stage.confidence_head = DSparkConfidenceHead(hidden_size=2, block_size=2) + model = nn.Module() + model.mtp_layers = nn.ModuleList([stage]) + model.config = SimpleNamespace( + q_lora_rank=1, + num_attention_heads=1, + num_hidden_layers=0, + num_nextn_predict_layers=1, + kv_lora_rank=1, + ) + model.model_config = SimpleNamespace( + mapping=SimpleNamespace( + tp_rank=0, + tp_size=1, + cp_rank=0, + cp_size=1, + enable_attention_dp=False, + ) + ) + model.post_load_weights = Mock() + model.cache_attn_weights_from_state_dict = Mock() + wrapper.dspark_model = model + wrapper.config = model.config + wrapper.num_stages = 1 + # Remapping has separate checkpoint-layout tests. This test starts with + # model-named tensors and does not mock the loader or its parameter copies. + remapper = "remap_dspark_v41_draft_keys" if v41 else "remap_dspark_draft_keys" + monkeypatch.setattr(module, remapper, lambda weights, *args, **kwargs: weights) + weights = {name: torch.full_like(param, 2.0) for name, param in model.named_parameters()} + return wrapper, weights + + +@pytest.mark.parametrize("v41", [False, True]) +@pytest.mark.parametrize("bias", [False, True]) +def test_embedded_confidence_load_validates_and_copies_real_parameters(monkeypatch, v41, bias): + wrapper, weights = _wrapper(monkeypatch, v41) + head = wrapper.dspark_model.mtp_layers[-1].confidence_head + prefix = "mtp_layers.0.confidence_head." + if bias: + weights[prefix + "proj.bias"] = torch.tensor([3.0]) + weight_pointer = head.proj.weight.data_ptr() + rng_before = torch.get_rng_state() + wrapper.load_weights(weights) + assert head.proj.weight.data_ptr() == weight_pointer + assert torch.equal(torch.get_rng_state(), rng_before) + assert torch.equal(head.proj.weight, weights[prefix + "proj.weight"]) + assert (head.proj.bias is not None) is bias + if bias: + assert torch.equal(head.proj.bias, weights[prefix + "proj.bias"]) + wrapper.dspark_model.post_load_weights.assert_called_once() + + +@pytest.mark.parametrize("v41", [False, True]) +@pytest.mark.parametrize("fault", ["unknown", "missing_weight", "broadcast_weight"]) +def test_embedded_confidence_rejects_incomplete_or_unknown_tensors(monkeypatch, v41, fault): + wrapper, weights = _wrapper(monkeypatch, v41) + prefix = "mtp_layers.0.confidence_head." + if fault == "unknown": + weights[prefix + "bias"] = torch.tensor([99.0]) + elif fault == "missing_weight": + weights.pop(prefix + "proj.weight") + weights[prefix + "proj.bias"] = torch.tensor([3.0]) + else: + weights[prefix + "proj.weight"] = torch.tensor([[2.0]]) + with pytest.raises(ValueError, match="confidence head checkpoint"): + wrapper.load_weights(weights) + wrapper.dspark_model.post_load_weights.assert_not_called() + + +def test_v41_entirely_absent_confidence_head_remains_supported(monkeypatch): + wrapper, weights = _wrapper(monkeypatch, True) + weights.pop("mtp_layers.0.confidence_head.proj.weight") + wrapper.load_weights(weights) + assert wrapper.dspark_model.mtp_layers[-1].confidence_head is None + wrapper.dspark_model.post_load_weights.assert_called_once() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_executor_policy.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_executor_policy.py new file mode 100644 index 000000000000..f7d80b2e0110 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_executor_policy.py @@ -0,0 +1,326 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only tests for DSpark policy agreement in the executor.""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ADPShapeAgreement +from tensorrt_llm._torch.pyexecutor.py_executor import ( + _DSPARK_ADP_WIRE_TRAILER_LEN, + _DSPARK_EXACT_NATIVE_YIELD_INDEX, + _DSPARK_EXACT_WIRE_PREFIX_LEN, + PyExecutor, + _classify_dspark_exact_generation_rows, + _decode_dspark_exact_expected_yields, + _dspark_exact_common_graph_batch_size, + _dspark_exact_secondary_padding_ready, + _encode_dspark_exact_expected_yields, + _validate_dspark_adp_acceptance_gate, + _validate_dspark_exact_bucket, +) + + +class _Batch: + def __init__(self, generation_requests): + self.generation_requests = generation_requests + self.context_requests = [] + + @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 _request(*, adp_dummy=False, other_dummy=False): + return SimpleNamespace( + is_dummy=bool(adp_dummy or other_dummy), + is_attention_dp_dummy=bool(adp_dummy), + ) + + +def _payload(*, rows, native, compact): + payload = [0] * (_DSPARK_EXACT_WIRE_PREFIX_LEN + len(compact) + _DSPARK_ADP_WIRE_TRAILER_LEN) + payload[1] = rows + payload[_DSPARK_EXACT_NATIVE_YIELD_INDEX] = native + payload[_DSPARK_EXACT_WIRE_PREFIX_LEN : _DSPARK_EXACT_WIRE_PREFIX_LEN + len(compact)] = compact + return payload + + +def _agreement(batch, peer_sizes=(3, 4)): + return ADPShapeAgreement( + iteration=9, + 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=True, + ) + + +def test_exact_yield_encoding_is_conservative_and_pads_missing_cells(): + exact_local = SimpleNamespace( + native_expected_yield=10.0, + compact_expected_yields=(9.25, float("nan")), + ) + + native, compact = _encode_dspark_exact_expected_yields( + exact_local=exact_local, num_cells=3, yield_scale=100 + ) + + assert native == 1000 + assert compact == (925, -1, -1) + + +def test_exact_yield_encoding_rejects_an_invalid_scale(): + with pytest.raises(ValueError, match="scale must be positive"): + _encode_dspark_exact_expected_yields(exact_local=None, num_cells=1, yield_scale=0) + + +def test_exact_yield_decode_aggregates_yield_and_worst_rank_loss(): + cells = ((128, 512), (128, 640)) + payloads = [ + _payload( + rows=2, + native=10_000_000, + compact=(8_000_000, 9_000_000), + ), + _payload( + rows=4, + native=20_000_000, + compact=(18_000_000, 14_000_000), + ), + _payload(rows=0, native=0, compact=(0, 0)), + ] + + native, compact, max_loss = _decode_dspark_exact_expected_yields( + payloads=payloads, + exact_cells=cells, + graph_batch_size=128, + yield_scale=1_000_000, + ) + + assert native == 30.0 + assert compact == {512: 26.0, 640: 23.0} + assert max_loss == {512: 1.0, 640: 1.5} + + +def test_exact_yield_decode_fails_closed_on_an_infeasible_peer(): + payloads = [ + _payload(rows=2, native=10_000_000, compact=(8_000_000,)), + _payload(rows=4, native=20_000_000, compact=(-1,)), + ] + + _, compact, max_loss = _decode_dspark_exact_expected_yields( + payloads=payloads, + exact_cells=((128, 512),), + graph_batch_size=128, + yield_scale=1_000_000, + ) + + assert compact == {512: 0.0} + assert max_loss == {512: 0.0} + + +def test_exact_row_classifier_authenticates_only_the_idle_adp_dummy(): + real = [_request(), _request()] + assert _classify_dspark_exact_generation_rows(real) == (real, False) + assert _classify_dspark_exact_generation_rows([]) == (None, False) + assert _classify_dspark_exact_generation_rows([_request(adp_dummy=True)]) == ( + [], + True, + ) + assert _classify_dspark_exact_generation_rows([real[0], _request(adp_dummy=True)]) == ( + None, + False, + ) + assert _classify_dspark_exact_generation_rows([_request(other_dummy=True)]) == (None, False) + + +def test_exact_common_graph_size_preserves_all_rank_termination(): + def round_up(rows): + return 16 if rows <= 16 else 32 + + assert _dspark_exact_common_graph_batch_size([[0, 16], [0, 0]], round_up) == 16 + + def reject_round_up(_rows): + raise AssertionError("all-zero peers must not perform a graph lookup") + + assert _dspark_exact_common_graph_batch_size([[0, 0], [0, 0]], reject_round_up) == 0 + + +@pytest.mark.parametrize( + ("budget", "secondary_ready", "expected"), + [(48, False, True), (49, False, True), (50, False, False), (50, True, True)], +) +def test_secondary_dummy_is_required_only_for_a_multirow_remainder( + budget, secondary_ready, expected +): + payloads = [[0, 16, 0], [0, 0, int(secondary_ready)]] + assert ( + _dspark_exact_secondary_padding_ready( + payloads, + graph_batch_size=16, + verifier_budget=budget, + secondary_ready_index=2, + ) + is expected + ) + + +def test_exact_bucket_must_match_the_selected_graph_key(): + _validate_dspark_exact_bucket(exact_shape=(16, 48, 3), bucket=48) + _validate_dspark_exact_bucket(exact_shape=None, bucket=None) + with pytest.raises(RuntimeError, match="selected captured verifier bucket"): + _validate_dspark_exact_bucket(exact_shape=(16, 48, 3), bucket=47) + + +def test_attention_dp_rejects_a_rank_local_acceptance_gate(): + with pytest.raises(ValueError, match="acceptance_rate_window_size"): + _validate_dspark_adp_acceptance_gate( + confidence_enabled=True, + attention_dp_enabled=True, + speculation_gate_enabled=True, + ) + _validate_dspark_adp_acceptance_gate( + confidence_enabled=True, + attention_dp_enabled=False, + speculation_gate_enabled=True, + ) + + +@pytest.mark.parametrize( + ("peer_sizes", "expected"), + [((3, 4), (True, True)), ((3, 0), (False, True))], +) +def test_can_queue_reuses_policy_agreement_without_another_collective(peer_sizes, expected): + batch = _Batch([object()] * 3) + + def reject_allgather(_value): + raise AssertionError("cached agreement issued another collective") + + executor = PyExecutor.__new__(PyExecutor) + executor.enable_attention_dp = True + executor.dist = SimpleNamespace(tp_size=8, tp_allgather=reject_allgather) + executor.model_engine = SimpleNamespace( + _dspark_confidence_enabled=True, + cuda_graph_runner=SimpleNamespace(adp_shape_agreement=_agreement(batch, peer_sizes)), + ) + executor.kv_connector_manager = None + executor.kv_cache_transceiver = None + executor.drafter = None + executor.iter_counter = 9 + executor._dspark_dynamic_handled_signature = (9, id(batch), 3, 0, 3) + executor._handle_dynamic_draft_len = lambda _batch: None + + assert executor._can_queue(batch) == expected + + +def _draft_len_executor(*, confidence_enabled): + spec_mode = SimpleNamespace( + support_dynamic_draft_len=lambda: False, + is_pard=lambda: False, + ) + spec_config = SimpleNamespace( + draft_len_schedule=None, + spec_dec_mode=spec_mode, + enable_confidence_scheduling=confidence_enabled, + use_rejection_sampling=False, + is_linear_tree=True, + ) + executor = PyExecutor.__new__(PyExecutor) + executor.model_engine = SimpleNamespace( + max_draft_len=5, + max_total_draft_tokens=5, + runtime_draft_len=None, + spec_config=spec_config, + cuda_graph_runner=SimpleNamespace(adp_shape_agreement=None), + ) + executor.speculation_permanently_disabled = False + executor.iter_counter = 9 + return executor + + +def test_fixed_draft_len_keeps_static_path_without_dynamic_signature(): + executor = _draft_len_executor(confidence_enabled=False) + batch = _Batch([object()]) + + executor._handle_dynamic_draft_len(batch) + executor._handle_dynamic_draft_len(batch) + + assert executor.model_engine.runtime_draft_len == 5 + assert not hasattr(executor, "_dspark_dynamic_handled_signature") + + +def test_confidence_draft_len_runs_only_once_per_iteration(): + executor = _draft_len_executor(confidence_enabled=True) + request = SimpleNamespace( + py_draft_tokens=[1, 2, 3, 4, 5], + py_needs_onehot_draft_probs=False, + ) + batch = _Batch([request]) + calls = 0 + + def select_draft_len(_batch): + nonlocal calls + calls += 1 + return 3 + + executor._dspark_confidence_draft_len = select_draft_len + executor._handle_dynamic_draft_len(batch) + executor._handle_dynamic_draft_len(batch) + + assert calls == 1 + assert executor.model_engine.runtime_draft_len == 3 + assert request.py_draft_tokens == [1, 2, 3] + assert executor._dspark_dynamic_handled_signature == (9, id(batch), 1, 0, 1) + + +def test_rebalance_suspends_both_dummy_variants_and_releases_the_pair_once(): + low = SimpleNamespace(py_request_id=10) + high = SimpleNamespace(py_request_id=11) + suspended_ids = [] + resumed_ids = [] + + class _Manager: + def is_request_active(self, request_id): + return request_id in {10, 11} + + def suspend_request(self, request): + suspended_ids.append(request.py_request_id) + + def resume_request(self, request): + resumed_ids.append(request.py_request_id) + return request.py_request_id != 11 + + released = [] + runner = SimpleNamespace( + padding_dummy_requests={5: low}, + secondary_padding_dummy_requests={5: high}, + release_padding_dummy=lambda resource, draft_len: released.append((resource, draft_len)), + ) + executor = PyExecutor.__new__(PyExecutor) + executor.model_engine = SimpleNamespace(cuda_graph_runner=runner) + executor.resource_manager = object() + manager = _Manager() + + suspended = executor._suspend_padding_dummies_for_rebalance(manager) + executor._resume_padding_dummies_after_rebalance(manager, suspended) + + assert suspended == [(5, 0, low), (5, 1, high)] + assert suspended_ids == [10, 11] + assert resumed_ids == [10, 11] + assert released == [(executor.resource_manager, 5)] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_fused_attention_route.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_fused_attention_route.py new file mode 100644 index 000000000000..7b1a6e030a4b --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_fused_attention_route.py @@ -0,0 +1,328 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU-only DSpark routing contracts; no Torch, CuTe, or CUDA runtime is loaded. + +Execute the production Python guards with tensor metadata and mocked launchers. +The AST extraction follows the existing K2/K3 routing-test prototype. Scalar +mask checks are not a substitute for compiling and testing the GPU kernels. +""" + +import ast +import math +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +_ROOT = Path(__file__).resolve().parents[5] +_OPS = _ROOT / "tensorrt_llm/_torch/custom_ops" +_ATTENTION = _OPS / "dspark_attention_custom_op.py" +_PREPARATION = _OPS / "dspark_rmsnorm_rope_custom_op.py" +_KERNELS = _ROOT / "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark" +_MODEL = _ROOT / "tensorrt_llm/_torch/models/modeling_dspark.py" +_PHYSICAL_KS = tuple(range(1, 9)) + + +def _tree(path): + return ast.parse(path.read_text(encoding="utf-8")) + + +def _load(path, name, namespace, owner=None): + nodes = _tree(path).body + if owner is not None: + nodes = next(node for node in nodes if getattr(node, "name", None) == owner).body + node = next(node for node in nodes if getattr(node, "name", None) == name) + node.decorator_list = [] + future = ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0) + module = ast.fix_missing_locations(ast.Module(body=[future, node], type_ignores=[])) + exec(compile(module, str(path), "exec"), namespace) + return namespace[name] + + +def _constants(path): + return { + node.targets[0].id: ast.literal_eval(node.value) + for node in _tree(path).body + if isinstance(node, ast.Assign) + and isinstance(node.targets[0], ast.Name) + and node.targets[0].id.startswith("_DSV4_DSPARK_") + } + + +def _tensor(shape, dtype="bf16", *, cuda=True, contiguous=True, strides=None): + if strides is None: + strides = tuple(math.prod(shape[i + 1 :]) for i in range(len(shape))) + return SimpleNamespace( + shape=shape, + ndim=len(shape), + dtype=dtype, + is_cuda=cuda, + is_contiguous=lambda: contiguous, + numel=lambda: math.prod(shape), + stride=lambda dim=None: strides if dim is None else strides[dim], + data_ptr=lambda: 0, + ) + + +def _guards(sm=100): + namespace = { + "torch": SimpleNamespace(bfloat16="bf16", float32="fp32", int32="i32", int64="i64"), + "get_sm_version": lambda: sm, + "_log_unsupported": lambda *args: False, + } + namespace.update(_constants(_PREPARATION)) + for name in ( + "_get_dspark_arch_str", + "is_fused_dspark_rmsnorm_rope_supported", + "_is_dspark_cache_write_supported", + "_is_dspark_draft_block_supported", + "is_fused_dspark_attention_preparation_supported", + ): + _load(_PREPARATION, name, namespace) + namespace.update(_constants(_ATTENTION)) + for name in ( + "is_dsv4_dspark_attention_config_supported", + "is_fused_dsv4_dspark_attention_supported", + "_is_fused_dsv4_dspark_attention_input_supported", + ): + _load(_ATTENTION, name, namespace) + return namespace + + +def _inputs(k, batch=2): + return [ + _tensor((batch, k, 128, 512)), + _tensor((batch, 8, 512)), + _tensor((batch + 1, 128, 512), contiguous=False, strides=(196608, 512, 1)), + _tensor((batch,), "i32"), + _tensor((batch,), "i32"), + _tensor((batch,), "i64"), + _tensor((128,), "fp32"), + _tensor((batch, k, 32, 2), "fp32"), + ] + + +class TestDSparkFusedAttentionRoute(unittest.TestCase): + def test_storage_bounded_range_is_consistent(self): + for path in (_ATTENTION, _PREPARATION): + constants = _constants(path) + self.assertEqual(constants["_DSV4_DSPARK_BLOCK_SIZES"], _PHYSICAL_KS) + self.assertEqual(constants["_DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE"], 8) + kernel = next( + node + for node in _tree(_KERNELS / "attention.py").body + if isinstance(node, ast.ClassDef) and node.name == "DSparkAttention" + ) + page_size = next( + ast.literal_eval(node.value) + for node in kernel.body + if isinstance(node, ast.Assign) and node.targets[0].id == "page_size_draft" + ) + self.assertEqual(_PHYSICAL_KS, tuple(range(1, page_size + 1))) + + def test_config_routes_every_k_and_rejects_out_of_range(self): + supported = _guards()["is_dsv4_dspark_attention_config_supported"] + for k in (-1, 0, *_PHYSICAL_KS, 9, 16, 128): + with self.subTest(k=k): + self.assertEqual(supported(k, 128, 512, 128), k in _PHYSICAL_KS) + for k in _PHYSICAL_KS: + self.assertFalse(supported(k, 64, 512, 128)) + self.assertFalse(supported(k, 128, 256, 128)) + self.assertFalse(supported(k, 128, 512, 64)) + + def test_kernel_constructor_keeps_geometry_and_storage_guards(self): + class KernelBase: + def __init__(self, *args, **kwargs): + self.seq_len_q = kwargs["seq_len_q"] + + namespace = { + "cutlass": SimpleNamespace(Float32="fp32"), + "cute": SimpleNamespace(jit=lambda f: f), + "DSparkAttentionKernel": KernelBase, + } + kernel = _load(_KERNELS / "attention.py", "DSparkAttention", namespace) + args = ("fp32", (128, 128), (128, 256), 1, 8, 128, 0.0) + for k in _PHYSICAL_KS: + with self.subTest(k=k): + op = kernel(*args, arch_str="sm_100", seq_len_q=k, inverse_rope_dim=64) + self.assertEqual(op.seq_len_q, k) + self.assertEqual(op.fixed_cache_seq_len, 128 + k) + self.assertEqual(op.tma_page_size_draft, 128) + for k in (-1, 0, 9, 16): + with self.subTest(k=k), self.assertRaises(ValueError): + kernel(*args, arch_str="sm_100", seq_len_q=k) + with self.assertRaises(ValueError): + kernel(*args, arch_str="sm_100", seq_len_q=8, inverse_rope_dim=32) + with self.assertRaises(ValueError): + kernel(*args[:4], 16, *args[5:], arch_str="sm_100", seq_len_q=8) + + def test_attention_tensor_contract_for_every_k_and_architecture(self): + for sm in (90, 100, 101, 103, 109, 120, 121): + gate = _guards(sm)["_is_fused_dsv4_dspark_attention_input_supported"] + for k in (0, *_PHYSICAL_KS, 9, 16): + with self.subTest(sm=sm, k=k): + self.assertEqual(gate(*_inputs(k)), sm in (100, 103) and k in _PHYSICAL_KS) + + def test_attention_dtype_shape_and_layout_guards_survive(self): + gate = _guards()["_is_fused_dsv4_dspark_attention_input_supported"] + bad_inputs = ( + (0, _tensor((2, 4, 128, 512), dtype="fp32")), + (0, _tensor((2, 4, 64, 512))), + (0, _tensor((2, 4, 128, 256))), + (0, _tensor((2, 4, 512))), + (1, _tensor((2, 7, 512))), + (1, _tensor((2, 9, 512))), + (1, _tensor((2, 8, 512), dtype="fp32")), + (2, _tensor((3, 64, 512))), + (2, _tensor((3, 128, 512), dtype="fp32")), + (2, _tensor((3, 128, 512), strides=(0, 512, 1))), + (2, _tensor((3, 128, 512), strides=(196609, 512, 1))), + (2, _tensor((3, 128, 512), strides=(196608, 513, 1))), + (2, _tensor((3, 128, 512), strides=(196608, 512, 2))), + (3, _tensor((2,), "i64")), + (4, _tensor((2,), "i64")), + (4, _tensor((3,), "i32")), + (5, _tensor((2,), "i32")), + (5, _tensor((3,), "i64")), + (6, _tensor((128,), "bf16")), + (6, _tensor((64,), "fp32")), + (7, _tensor((2, 4, 31, 2), "fp32")), + ) + for index, bad in bad_inputs: + with self.subTest(index=index, shape=bad.shape, dtype=bad.dtype): + inputs = _inputs(4) + inputs[index] = bad + self.assertFalse(gate(*inputs)) + for k in _PHYSICAL_KS: + for index in range(8): + inputs = _inputs(k) + inputs[index].is_cuda = False + self.assertFalse(gate(*inputs)) + if index != 2: # The worker's strided cache-page view is supported. + inputs = _inputs(k) + inputs[index].is_contiguous = lambda: False + self.assertFalse(gate(*inputs)) + + def test_preparation_bounds_and_guards_for_every_k(self): + gate = _guards()["is_fused_dspark_attention_preparation_supported"] + for k in (0, *_PHYSICAL_KS, 9, 16): + with self.subTest(k=k): + inputs = [ + _tensor((2, 1, 512)), + _tensor((2, k, 512)), + _tensor((512,)), + _tensor((2, 32, 2), "fp32"), + _tensor((2 * k, 32, 2), "fp32"), + _inputs(k)[2], + _tensor((2,), "i64"), + _tensor((2,), "i64"), + ] + self.assertEqual(gate(*inputs), k in _PHYSICAL_KS) + if k in _PHYSICAL_KS: + for index in range(len(inputs)): + original = inputs[index].is_cuda + inputs[index].is_cuda = False + self.assertFalse(gate(*inputs)) + inputs[index].is_cuda = original + inputs[4] = _tensor((2 * k + 1, 32, 2), "fp32") + self.assertFalse(gate(*inputs)) + + def test_masks_cover_short_full_wrapped_and_empty_windows(self): + namespace = { + "cutlass": SimpleNamespace(Int32=int, const_expr=lambda x: x), + "cute": SimpleNamespace(elem_less=lambda a, b: a < b), + } + path = _KERNELS / "attention_kernel.py" + mask = _load(path, "is_score_valid", namespace, "DSparkAttentionKernel") + valid_length = _load(path, "get_window_valid_len", namespace, "DSparkAttentionKernel") + kernel = SimpleNamespace( + mma_qk_tiler=(128, 128, 128), + window_valid_len_from_tensor=True, + tma_page_size_win=128, + ) + for k in _PHYSICAL_KS: + for end in (0, 5, 127, 128, 257, 390): + for valid in (-1, 0, 1, 3, 128, 129): + with self.subTest(k=k, end=end, valid=valid): + length = valid_length(kernel, [valid], 0) + expected = {(end - age) % 128 for age in range(min(max(valid, 0), 128))} + actual = { + col for col in range(128) if mask(kernel, col, 0, 128 + k, length, end) + } + self.assertEqual(actual, expected) + draft = { + col for col in range(128) if mask(kernel, col, 1, 128 + k, length, end) + } + self.assertEqual(draft, set(range(k))) + + def test_cache_keys_separate_k_and_arch_but_not_runtime_shapes(self): + namespace = _guards() + cache = {} + compile_kernel = Mock(side_effect=lambda *args: Mock()) + namespace["_dspark_attention_kernel_cache"] = cache + namespace["_compile_dspark_attention"] = compile_kernel + capturing = Mock(return_value=False) + namespace["torch"].cuda = SimpleNamespace(is_current_stream_capturing=capturing) + namespace["torch"].empty_like = lambda q: _tensor(q.shape) + run = _load(_ATTENTION, "_run_dspark_attention", namespace) + for arch in ("sm_100", "sm_103"): + namespace["_get_dspark_arch_str"] = lambda arch=arch: arch + for k in _PHYSICAL_KS: + with self.subTest(k=k, arch=arch): + for batch, scale in ((1, 0.125), (3, 0.25)): + self.assertEqual(run(*_inputs(k, batch), scale).shape, (batch, k, 128, 512)) + self.assertEqual(cache[(k, arch)].call_count, 2) + self.assertEqual(compile_kernel.call_count, 16) + self.assertEqual(capturing.call_count, 16) + self.assertEqual(len({id(compiled) for compiled in cache.values()}), 16) + + def test_graph_capture_only_uses_warmed_k(self): + for k in _PHYSICAL_KS: + namespace = _guards() + namespace["_dspark_attention_kernel_cache"] = {} + namespace["_compile_dspark_attention"] = Mock(side_effect=AssertionError("compile")) + namespace["torch"].cuda = SimpleNamespace(is_current_stream_capturing=lambda: True) + namespace["torch"].empty_like = Mock(side_effect=AssertionError("allocation")) + run = _load(_ATTENTION, "_run_dspark_attention", namespace) + with self.subTest(k=k), self.assertRaisesRegex(RuntimeError, "warmed up"): + run(*_inputs(k), 0.125) + + def test_unsupported_k_warmup_does_not_allocate(self): + namespace = _guards() + warmup = _load(_ATTENTION, "warmup_fused_dsv4_dspark_attention", namespace) + # The metadata-only torch object cannot allocate; early return is required. + for k in (-1, 0, 9, 16): + warmup(k, 1e-6) + + def test_query_storage_and_fallback_routes_remain_separate(self): + attention = ast.unparse(_tree(_KERNELS / "attention.py")) + self.assertIn("(self.num_heads, self.head_dim, self.seq_len_q, batch_size)", attention) + self.assertIn("batch_size * self.seq_len_q * self.inverse_rope_dim", attention) + self.assertIn("self.fixed_cache_seq_len = self.window_size + seq_len_q", attention) + model = {} + _load(_MODEL, "dspark_attention_forward_batched", model) + function = next( + node + for node in _tree(_MODEL).body + if getattr(node, "name", None) == "dspark_attention_forward_batched" + ) + route = next( + node + for node in function.body + if isinstance(node, ast.If) + and ast.unparse(node.test) == "use_fused_dsv4_dspark_attention" + ) + fused = ast.unparse(ast.Module(body=route.body, type_ignores=[])) + fallback = ast.unparse(ast.Module(body=route.orelse, type_ignores=[])) + self.assertIn("cute_dsl_dspark_rmsnorm_rope_draft_block(", fused) + self.assertIn("fused_dsv4_dspark_attention(", fused) + self.assertNotIn("fused_dsv4_dspark_attention(", fallback) + self.assertIn( + "get_dspark_topk_idxs_batched(window_size, block, start_pos, valid_len)", fallback + ) + self.assertIn("torch.cat([cache_rows, kv], dim=1)", fallback) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_evidence.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_evidence.py new file mode 100644 index 000000000000..3fca02f00a02 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_evidence.py @@ -0,0 +1,458 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only, deliberately synthetic protocol fixtures; no production evidence.""" + +import copy +import dataclasses +import hashlib +import json +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace + +from tensorrt_llm._torch.speculative import dspark_mixed_evidence as mixed +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig, DSparkMixedEvidenceConfig + + +def sha(raw): + return hashlib.sha256(raw).hexdigest() + + +def encode(value): + return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() + + +def identity(): + return mixed.MixedRuntimeIdentity(*("abcdef"[index] * 64 for index in range(6))) + + +def geometry(): + # Rank 0 prefills 64 rows; rank 1 decodes with native 32 versus compact 16. + # Every bootstrap/dummy generation is still charged its full K+1 rows. + features = ( + (64, 4096, 8192, 0, 0, 0, 0, 1, 0, 1, 0, 0), + (0, 0, 0, 0, 1024, 128, 0, 0, 8, 0, 8, 0), + ) + ((0,) * 12,) * 6 + return mixed.MixedGeometry(features, 96, (64, 16, 0, 0, 0, 0, 0, 0), (64, 32, 0, 0, 0, 0, 0, 0)) + + +def fixture(): + live = identity() + table = { + "schema": "dspark-mixed-breakable-costs-v1", + "identity": live.payload(), + "protocol_sha256": "1" * 64, + "action_semantics": mixed._ACTION, + "pricing_semantics": mixed._PRICE_SEMANTICS, + "synthetic": False, + "provenance": "closed_measured_breakable_mixed", + "minimum_predicted_gain": 0.01, + "cells": [ + { + "geometry": geometry().payload(), + "native_lower_ms": 7.0, + "candidate_upper_ms": 6.9, + "measurement_result_sha256": "2" * 64, + } + ], + } + receipt = { + "schema": "dspark-mixed-breakable-admission-v1", + "identity": live.payload(), + "costs_sha256": sha(encode(table)), + "protocol_sha256": "1" * 64, + "action_semantics": mixed._ACTION, + "pricing_semantics": mixed._PRICE_SEMANTICS, + "production_admission": True, + "source_and_measurements_verified": True, + "accuracy_pass": True, + "runtime_invariants_pass": True, + "matched_workload_pass": True, + "accuracy_result_sha256": "3" * 64, + "runtime_invariants_result_sha256": "4" * 64, + "matched_workload_result_sha256": "5" * 64, + "selector_replay_sha256": "6" * 64, + "measurement_results_sha256": sha(encode(["2" * 64])), + "economics": { + "policy_steps": 10, + "compact_choices": 4, + "gross_compact_value_ms_lower_bound": 3.0, + "fixed_confidence_path_overhead_ms_upper_bound": 0.1, + "safety_margin_ms": 0.5, + }, + } + return table, receipt, live + + +def resolve(table, receipt, live=None, *, update_receipt=True): + table_raw = encode(table) + receipt = copy.deepcopy(receipt) + if update_receipt: + receipt["costs_sha256"] = sha(table_raw) + receipt_raw = encode(receipt) + binding = mixed.MixedEvidenceBinding(sha(table_raw), sha(receipt_raw), "1" * 64) + return mixed.resolve_mixed_evidence( + binding, costs_raw=table_raw, admission_raw=receipt_raw, live_identity=live or identity() + ) + + +class TestMixedEvidence(unittest.TestCase): + def assert_native(self, table, receipt, live=None, **kwargs): + result = resolve(table, receipt, live, **kwargs) + self.assertIsNone(result.costs) + self.assertTrue(result.reason.startswith("mixed_evidence_invalid_native:")) + + def test_01_actual_binding_positive_synthetic_fixture_only(self): + table, receipt, live = fixture() + resolved = resolve(table, receipt, live) + self.assertEqual(resolved.reason, "mixed_evidence_bound") + self.assertEqual(resolved.costs.lookup(geometry()).candidate_upper_ms, 6.9) + self.assertEqual(resolved.costs.minimum_predicted_gain, 0.01) + + def test_02_bad_fixed_restored_old_validation_schema(self): + table, receipt, live = fixture() + table["schema"] = "dspark-host-ragged-bcg-prices-v1" + self.assert_native(table, receipt, live) + table["schema"] = "dspark-mixed-breakable-costs-v1" + self.assertIsNotNone(resolve(table, receipt, live).costs) + table["schema"] = "dspark-host-ragged-bcg-prices-v1" + self.assert_native(table, receipt, live) + + def test_03_absent_optional_binding_does_not_read_payloads(self): + result = mixed.resolve_mixed_evidence( + None, costs_raw=None, admission_raw=None, live_identity=identity() + ) + self.assertEqual(result.reason, "mixed_evidence_absent_native") + self.assertIsNone(result.costs) + + def test_04_decode_policy_object_and_flag_unchanged_for_all_outcomes(self): + decode = {"enabled": True, "identity": "existing-decode", "table": object()} + before = dict(decode) + table, receipt, live = fixture() + for binding in (None, mixed.MixedEvidenceBinding("0" * 64, "1" * 64, "2" * 64)): + mixed.resolve_mixed_evidence( + binding, costs_raw=b"bad", admission_raw=b"bad", live_identity=live + ) + self.assertEqual(decode, before) + self.assertTrue(decode["enabled"]) + self.assertIsNotNone(resolve(table, receipt, live).costs) + self.assertEqual(decode, before) + + def test_05_live_identity_is_not_read_from_table(self): + table, receipt, live = fixture() + for field in ( + "source_tree_sha256", + "runtime_snapshot_sha256", + "model_config_sha256", + "gpu_class_sha256", + "topology_sha256", + "selector_identity_sha256", + ): + with self.subTest(field=field): + self.assert_native(table, receipt, dataclasses.replace(live, **{field: "9" * 64})) + + def test_06_invalid_supported_runtime_and_mutable_grid_rejected(self): + live = identity() + for overrides in ( + {"physical_k": 2}, + {"physical_k": True}, + {"world_size": 1}, + {"execution_mode": "FULL"}, + {"body_buckets": list(live.body_buckets)}, + {"body_buckets": (96, 128, 192, 256)}, + ): + with self.subTest(overrides=overrides), self.assertRaises(ValueError): + dataclasses.replace(live, **overrides) + + def test_07_payloads_are_deeply_frozen_into_typed_tuples(self): + table, receipt, live = fixture() + costs = resolve(table, receipt, live).costs + table["cells"][0]["geometry"]["rank_features"][1]["eligible_generations"] = 0 + table["cells"][0]["native_lower_ms"] = 0 + self.assertEqual(costs.lookup(geometry()).native_lower_ms, 7.0) + with self.assertRaises(dataclasses.FrozenInstanceError): + costs.cells[0].candidate_upper_ms = 1 + + def test_08_uncovered_geometry_is_native_no_interpolation(self): + table, receipt, live = fixture() + costs = resolve(table, receipt, live).costs + unseen = dataclasses.replace(geometry(), actual_rank_rows=(64, 15, 0, 0, 0, 0, 0, 0)) + self.assertIsNone(costs.lookup(unseen)) + self.assertIsNone(costs.lookup(geometry().payload())) + + def test_09_duplicate_geometry_rejects_favorable_reselection(self): + table, receipt, _ = fixture() + table["cells"].append(copy.deepcopy(table["cells"][0])) + table["cells"][1]["candidate_upper_ms"] = 1 + self.assert_native(table, receipt) + + def test_10_cell_list_empty_or_unbounded_rejected(self): + table, receipt, _ = fixture() + for cells in ([], [table["cells"][0]] * 257): + bad = copy.deepcopy(table) + bad["cells"] = cells + self.assert_native(bad, receipt) + + def test_11_context_and_fixed_generation_rows_cannot_be_discounted(self): + original = geometry() + for change in ( + {"actual_rank_rows": (63, 16, 0, 0, 0, 0, 0, 0)}, + {"actual_rank_rows": (64, 7, 0, 0, 0, 0, 0, 0)}, + {"native_rank_rows": (64, 31, 0, 0, 0, 0, 0, 0)}, + ): + with self.subTest(change=change), self.assertRaises(ValueError): + dataclasses.replace(original, **change) + features = list(original.rank_features) + features[1] = (0, 0, 0, 0, 1024, 128, 1, 0, 7, 0, 8, 0) + with self.assertRaises(ValueError): + dataclasses.replace(original, rank_features=tuple(features)) + + def test_12_decode_only_and_noncompact_geometry_rejected(self): + original = geometry() + features = list(original.rank_features) + features[0] = (64, 4096, 8192, 0, 0, 0, 0, 0, 0, 1, 0, 0) + with self.assertRaises(ValueError): + dataclasses.replace(original, rank_features=tuple(features)) + with self.assertRaises(ValueError): + dataclasses.replace(original, actual_rank_rows=original.native_rank_rows) + + def test_13_rank_order_is_not_coalesced(self): + table, receipt, live = fixture() + costs = resolve(table, receipt, live).costs + original = geometry() + reordered = dataclasses.replace( + original, + rank_features=original.rank_features[::-1], + actual_rank_rows=original.actual_rank_rows[::-1], + native_rank_rows=original.native_rank_rows[::-1], + ) + self.assertIsNone(costs.lookup(reordered)) + + def test_14_unknown_null_and_validation_only_fields_reject(self): + table, receipt, _ = fixture() + for field, value in ( + ("validation_only", False), + ("pricing_sequence", [96]), + ("confidence_sps_table_path", "old.json"), + ("unknown", None), + ): + bad = copy.deepcopy(table) + bad[field] = value + self.assert_native(bad, receipt) + + def test_15_synthetic_prices_and_claim_only_receipts_reject(self): + table, receipt, _ = fixture() + for field in ( + "production_admission", + "source_and_measurements_verified", + "accuracy_pass", + "runtime_invariants_pass", + "matched_workload_pass", + ): + bad = copy.deepcopy(receipt) + bad[field] = False + self.assert_native(table, bad) + table["synthetic"] = True + self.assert_native(table, receipt) + + def test_16_proxy_is_never_relabelled_as_whole_step(self): + table, receipt, _ = fixture() + for key in ("table", "receipt"): + bad_table, bad_receipt = copy.deepcopy(table), copy.deepcopy(receipt) + (bad_table if key == "table" else bad_receipt)["pricing_semantics"] = ( + "qualified_group_whole_step" + ) + self.assert_native(bad_table, bad_receipt) + + def test_17_protocol_action_and_measurement_binding_reject_drift(self): + table, receipt, _ = fixture() + for key, value in ( + ("protocol_sha256", "9" * 64), + ("action_semantics", "uniform-L"), + ("measurement_results_sha256", "9" * 64), + ): + bad = copy.deepcopy(receipt) + bad[key] = value + self.assert_native(table, bad) + + def test_18_all_result_hashes_required(self): + table, receipt, _ = fixture() + for field in ( + "accuracy_result_sha256", + "runtime_invariants_result_sha256", + "matched_workload_result_sha256", + "selector_replay_sha256", + "measurement_results_sha256", + ): + bad = copy.deepcopy(receipt) + bad[field] = "not-a-hash" + self.assert_native(table, bad) + + def test_19_positive_economics_recomputed_no_boolean_override(self): + table, receipt, _ = fixture() + for change in ( + {"compact_choices": 0}, + {"gross_compact_value_ms_lower_bound": 1.0}, + {"policy_steps": True}, + {"safety_margin_ms": -1.0}, + {"compact_choices": 11}, + ): + bad = copy.deepcopy(receipt) + bad["economics"].update(change) + self.assert_native(table, bad) + + def test_20_no_relaxation_of_one_percent_immediate_gate(self): + table, receipt, _ = fixture() + for margin in (0.0, 0.009, True): + table["minimum_predicted_gain"] = margin + self.assert_native(table, receipt) + table["minimum_predicted_gain"] = 0.02 + self.assertEqual(resolve(table, receipt).costs.minimum_predicted_gain, 0.02) + + def test_21_nonpositive_and_nonfinite_costs_rejected(self): + table, receipt, _ = fixture() + for field in ("native_lower_ms", "candidate_upper_ms"): + for value in (0, -1, True, "7.0"): + bad = copy.deepcopy(table) + bad["cells"][0][field] = value + self.assert_native(bad, receipt) + raw = encode(table).replace(b'"candidate_upper_ms":6.9', b'"candidate_upper_ms":1e9999') + receipt["costs_sha256"] = sha(raw) + self.assert_raw_native(raw, encode(receipt)) + + def assert_raw_native(self, table_raw, receipt_raw): + binding = mixed.MixedEvidenceBinding(sha(table_raw), sha(receipt_raw), "1" * 64) + result = mixed.resolve_mixed_evidence( + binding, costs_raw=table_raw, admission_raw=receipt_raw, live_identity=identity() + ) + self.assertIsNone(result.costs) + + def test_22_duplicate_nonfinite_malformed_deep_json_rejected(self): + _, receipt, _ = fixture() + for raw in ( + b'{"schema":"x","schema":"y"}', + b'{"x":NaN}', + b'{"x":Infinity}', + b"{malformed", + b"\xff", + b"[" * 1500 + b"]" * 1500, + ): + self.assert_raw_native(raw, encode(receipt)) + + def test_23_missing_raw_oversized_and_digest_drift_rejected(self): + table, receipt, live = fixture() + raw, qraw = encode(table), encode(receipt) + binding = mixed.MixedEvidenceBinding(sha(raw), sha(qraw), "1" * 64) + for value in (None, raw + b" ", bytearray(raw), b"", b" " * (mixed._MAX_BYTES + 1)): + result = mixed.resolve_mixed_evidence( + binding, costs_raw=value, admission_raw=qraw, live_identity=live + ) + self.assertIsNone(result.costs) + + def test_24_table_receipt_cross_binding_rejected(self): + table, receipt, _ = fixture() + table["cells"][0]["candidate_upper_ms"] = 6.8 + self.assert_native(table, receipt, update_receipt=False) + + def test_25_rank_extents_strict_types_and_wire_limits(self): + original = geometry() + for value in (True, 1.0, -1, 2**63): + features = list(original.rank_features) + changed = list(features[0]) + changed[2] = value + features[0] = tuple(changed) + with self.subTest(value=value), self.assertRaises((TypeError, ValueError)): + dataclasses.replace(original, rank_features=tuple(features)) + + def test_26_decode_schema_is_not_mixed_evidence(self): + _, receipt, _ = fixture() + self.assert_raw_native(encode({"schema_version": 2}), encode(receipt)) + + def test_27_valid_bootstrap_reserves_full_k_plus_one_rows(self): + original = geometry() + features = list(original.rank_features) + features[1] = (0, 0, 0, 0, 1024, 128, 1, 0, 7, 0, 8, 4) + valid = dataclasses.replace( + original, rank_features=tuple(features), actual_rank_rows=(64, 20, 0, 0, 0, 0, 0, 0) + ) + self.assertEqual(valid.rank_features[1][-1], 4) + with self.assertRaises(ValueError): + dataclasses.replace(valid, actual_rank_rows=(64, 10, 0, 0, 0, 0, 0, 0)) + + def test_28_bad_fixed_restored_cost_pin_is_not_relabelled(self): + table, receipt, live = fixture() + raw, qraw = encode(table), encode(receipt) + bad_binding = mixed.MixedEvidenceBinding("0" * 64, sha(qraw), "1" * 64) + for binding, expected in ( + (bad_binding, False), + (mixed.MixedEvidenceBinding(sha(raw), sha(qraw), "1" * 64), True), + (bad_binding, False), + ): + result = mixed.resolve_mixed_evidence( + binding, costs_raw=raw, admission_raw=qraw, live_identity=live + ) + self.assertEqual(result.costs is not None, expected) + + def test_29_file_binding_requires_independent_live_fingerprint(self): + table, receipt, live = fixture() + table_raw, receipt_raw = encode(table), encode(receipt) + live_raw = encode( + {"schema": "dspark-mixed-live-fingerprint-v1", "identity": live.payload()} + ) + with tempfile.TemporaryDirectory() as directory: + paths = [Path(directory) / name for name in ("costs.json", "receipt.json", "live.json")] + for path, raw in zip(paths, (table_raw, receipt_raw, live_raw)): + path.write_bytes(raw) + config = SimpleNamespace( + costs_path=str(paths[0]), + costs_sha256=sha(table_raw), + admission_path=str(paths[1]), + admission_sha256=sha(receipt_raw), + protocol_sha256="1" * 64, + live_fingerprint_path=str(paths[2]), + live_fingerprint_sha256=sha(live_raw), + ) + self.assertIsNotNone(mixed.load_mixed_evidence(config).costs) + config.live_fingerprint_sha256 = "0" * 64 + self.assertIsNone(mixed.load_mixed_evidence(config).costs) + self.assertEqual(mixed.load_mixed_evidence(None).reason, "mixed_evidence_absent_native") + + def test_30_typed_config_rejects_invalid_pins_and_unknown_fields(self): + data = dict( + costs_path="costs.json", + costs_sha256="1" * 64, + admission_path="receipt.json", + admission_sha256="2" * 64, + protocol_sha256="3" * 64, + live_fingerprint_path="live.json", + live_fingerprint_sha256="4" * 64, + ) + self.assertEqual(DSparkMixedEvidenceConfig(**data).costs_path, "costs.json") + with self.assertRaises(ValueError): + DSparkMixedEvidenceConfig(**{**data, "costs_sha256": "invalid"}) + with self.assertRaises(ValueError): + DSparkMixedEvidenceConfig(**data, force=True) + + def test_31_feature_off_rejects_mixed_input_and_k1_clears_it(self): + config = SimpleNamespace( + enable_confidence_scheduling=False, + enable_fused_confidence_scheduler=False, + confidence_sts_path=None, + confidence_sps_table_path=None, + confidence_sps_live_fingerprint_path=None, + confidence_admission_receipt_path=None, + confidence_admission_receipt_sha256=None, + confidence_mixed_evidence=object(), + max_draft_len=3, + ) + with self.assertRaisesRegex(ValueError, "require enable_confidence_scheduling"): + DSparkDecodingConfig.validate_confidence_scheduling(config) + config.enable_confidence_scheduling, config.max_draft_len = True, 1 + self.assertIs(DSparkDecodingConfig.validate_confidence_scheduling(config), config) + self.assertFalse(config.enable_confidence_scheduling) + self.assertIsNone(config.confidence_mixed_evidence) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_runtime.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_runtime.py new file mode 100644 index 000000000000..73f8cb0aa618 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_runtime.py @@ -0,0 +1,350 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU contracts for optional DSpark mixed allocation; prices are synthetic.""" + +import copy +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor import dspark_mixed_runtime as runtime +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType +from tensorrt_llm._torch.speculative.dspark_mixed_evidence import ( + AdmittedMixedCosts, + MixedCostCell, + MixedEvidenceBinding, + MixedGeometry, + MixedRuntimeIdentity, +) +from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig +from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner + + +class _Event: + def __init__(self, ready=True): + self.ready = ready + self.records = 0 + + def query(self): + return self.ready + + def record(self): + self.records += 1 + + +def _request(request_id, *, context=False): + request = SimpleNamespace( + py_request_id=request_id, + py_seq_slot=request_id, + is_dummy=False, + py_batch_idx=0, + py_draft_tokens=[1, 2, 3], + py_verify_len=None, + state=SimpleNamespace(name="GENERATION_IN_PROGRESS"), + py_orig_prompt_len=128, + py_max_new_tokens=768, + context_chunk_size=64 if context else 0, + context_current_position=128 if context else 0, + py_num_compressed_tokens=0, + is_first_context_chunk=False, + ) + request.get_num_tokens = lambda beam: 128 + return request + + +def _group(monkeypatch, *, all_decode=False, stale=False): + monkeypatch.setattr(runtime, "_greedy", lambda _batch: True) + identity = MixedRuntimeIdentity(*(c * 64 for c in "abcdef")) + binding = MixedEvidenceBinding("1" * 64, "2" * 64, "3" * 64) + costs = AdmittedMixedCosts(identity, binding, 0.01, ()) + group = [] + for rank in range(8): + config = SimpleNamespace( + max_draft_len=3, + max_seq_len=4096, + max_num_tokens=4096, + enable_attention_dp=True, + torch_compile_enabled=False, + is_multimodal=False, + spec_config=SimpleNamespace( + max_draft_len=3, + block_size=3, + decoding_type="DSpark", + draft_is_embedded_in_target=True, + use_rejection_sampling=False, + ), + prefill_cuda_graph_backend=SimpleNamespace(name="BREAKABLE"), + prefill_cuda_graph_num_tokens=runtime.BODY_BUCKETS, + ) + generations = [_request(rank * 1000 + i + 1) for i in range(64)] + contexts = [_request(99999, context=True)] if rank == 0 and not all_decode else [] + batch = SimpleNamespace( + context_requests=contexts, + generation_requests=generations, + context_requests_last_chunk=contexts, + ) + planner = DSparkVerifyPlanner(cfg=DSparkScheduleConfig(3), device_windows=True) + planner._host_buffer = torch.full((64, 3), 5.0) + planner._host_confidence_stamp = torch.full((64,), 4, dtype=torch.int32) + planner._snapshot_valid, planner._copy_event = True, _Event() + slots = {r.py_request_id: i for i, r in enumerate(generations)} + incarnations = {r.py_request_id: i + 1 for i, r in enumerate(generations)} + worker = SimpleNamespace( + _draft_seq_host=5, + _req_to_slot=slots, + confidence_row_for=slots.__getitem__, + confidence_incarnation_for=incarnations.get, + ) + planner._mixed_current_snapshot_meta = dict( + buffer=planner._host_buffer, + event=planner._copy_event, + owners=slots.copy(), + staged_iteration=9, + staging_sequence=5, + confidence_stamp=planner._host_confidence_stamp, + producer_attempts={r.py_request_id: (r, i, i + 1) for i, r in enumerate(generations)}, + ) + if stale and rank == 1: + incarnations[generations[0].py_request_id] += 100 + runner = SimpleNamespace( + _config=config, + mapping=SimpleNamespace(tp_size=8, pp_size=1, cp_size=1), + _dspark_confidence_enabled=True, + _dspark_trims_submitted_tokens=True, + use_beam_search=False, + _dspark_device_budget=10, + cuda_graph_runner=SimpleNamespace( + agreed_ragged_bucket=192, + ragged_pad_verify_len=1, + ragged_zero_real_high_rows=0, + ), + breakable_cuda_graph_runner=SimpleNamespace( + is_capturing=False, is_warming_up=False, has_graph=lambda _bucket: True + ), + ) + group.append((runtime.MixedRuntime(runner, costs), batch, worker, planner, incarnations)) + votes = [ + mixed.local_vote(batch, worker, planner, 10, rank, executor_eligible=True) + for rank, (mixed, batch, worker, planner, _) in enumerate(group) + ] + if not all_decode: + index = runtime.BODY_BUCKETS.index(192) + geometry = MixedGeometry( + tuple(tuple(v[14:26]) for v in votes), + 192, + tuple(v[28 + index * 3 + 2] for v in votes), + tuple(v[27] for v in votes), + ) + costs = replace(costs, cells=(MixedCostCell(geometry, 10.0, 5.0, "4" * 64),)) + for mixed, *_ in group: + mixed.costs = costs + return group, votes + + +def test_host_mixed_shortens_decode_peer_without_mutating_context(monkeypatch): + group, votes = _group(monkeypatch) + context = copy.deepcopy(group[0][1].context_requests[0].__dict__) + for mixed, batch, worker, planner, _ in group: + assert mixed.apply(batch, worker, planner, votes, 10) + assert mixed.runner._dspark_host_window_step + assert mixed.runner.cuda_graph_runner._dspark_host_window_batch is batch + assert mixed.runner._dspark_device_budget is None + assert any(r.py_verify_len < 3 for r in batch.generation_requests) + assert group[0][1].context_requests[0].__dict__ == context + assert sum(r.py_verify_len + 1 for r in group[1][1].generation_requests) == 192 + + +@pytest.mark.parametrize("failure", ["uncovered", "unready", "identity"]) +def test_missing_mixed_evidence_stays_native_without_disabling_decode(monkeypatch, failure): + group, votes = _group(monkeypatch) + if failure == "uncovered": + for mixed, *_ in group: + mixed.costs = replace(mixed.costs, cells=()) + elif failure == "unready": + mixed, batch, worker, planner, _ = group[1] + planner._copy_event.ready = False + votes[1] = mixed.local_vote(batch, worker, planner, 10, 1, executor_eligible=True) + assert votes[1][5] == 0 + else: + votes[2][6] ^= 1 + for mixed, batch, worker, planner, _ in group: + assert mixed.apply(batch, worker, planner, votes, 10) + assert all(r.py_verify_len == 3 for r in batch.generation_requests) + assert mixed.runner._dspark_confidence_enabled + + +def test_stale_owner_keeps_full_width_without_discarding_other_owners(monkeypatch): + group, votes = _group(monkeypatch, stale=True) + mixed, batch, worker, planner, _ = group[1] + assert votes[1][22] == 63 + mixed.apply(batch, worker, planner, votes, 10) + assert batch.generation_requests[0].py_verify_len == 3 + assert any(r.py_verify_len < 3 for r in batch.generation_requests[1:]) + + +@pytest.mark.parametrize("change", ["incarnation", "context", "first_chunk"]) +def test_ticket_rechecked_before_publishing_lengths(monkeypatch, change): + group, votes = _group(monkeypatch) + rank = 1 if change == "incarnation" else 0 + mixed, batch, worker, planner, incarnations = group[rank] + if change == "incarnation": + incarnations[batch.generation_requests[0].py_request_id] += 1 + elif change == "context": + batch.context_requests[0].context_current_position += 1 + else: + batch.context_requests[0].is_first_context_chunk = True + with pytest.raises(ValueError, match="changed"): + mixed.apply(batch, worker, planner, votes, 10) + assert all(r.py_verify_len is None for r in batch.generation_requests) + + +@pytest.mark.parametrize("unresolved", ["first_chunk", "token_budget"]) +def test_unfinalized_geometry_declines_group_before_compact_publication(monkeypatch, unresolved): + group, votes = _group(monkeypatch) + mixed, batch, worker, planner, _ = group[0] + if unresolved == "first_chunk": + batch.context_requests[0].is_first_context_chunk = True + else: + mixed.runner._config.max_num_tokens = 319 + votes[0] = mixed.local_vote(batch, worker, planner, 10, 0, executor_eligible=True) + assert votes[0][5] == 0 + assert all(r.py_verify_len is None for r in batch.generation_requests) + for mixed, batch, worker, planner, _ in group: + assert mixed.apply(batch, worker, planner, votes, 10) + assert all(r.py_verify_len == 3 for r in batch.generation_requests) + assert mixed.runner._dspark_host_window_step + + +def test_first_chunk_stays_native_across_real_python_resource_preparation(monkeypatch): + group, votes = _group(monkeypatch) + mixed, batch, worker, planner, _ = group[0] + + class _NativeRequestModel(SimpleNamespace): + # Models the native first-chunk predicate; no C++ execution is claimed. + @property + def is_first_context_chunk(self): + return self.context_current_position == self.prepopulated_prompt_len + + values = batch.context_requests[0].__dict__.copy() + values.pop("is_first_context_chunk") + context = _NativeRequestModel(**values) + context.context_chunk_size, context.context_current_position = 256, 0 + context.prepopulated_prompt_len = 0 + batch.context_requests[:] = [context] + batch.context_requests_last_chunk[:] = [context] + votes[0] = mixed.local_vote(batch, worker, planner, 10, 0, executor_eligible=True) + assert votes[0][5] == 0 + for mixed, member, worker, planner, _ in group: + mixed.apply(member, worker, planner, votes, 10) + assert all(r.py_verify_len == 3 for r in member.generation_requests) + phases = [] + + def model_prefix_cache_setter(_batch): + phases.append("prepare") + context.prepopulated_prompt_len = context.context_current_position = 128 + context.context_chunk_size = 128 + + manager = SimpleNamespace( + resource_managers={ + ResourceManagerType.KV_CACHE_MANAGER: SimpleNamespace( + prepare_resources=model_prefix_cache_setter, + report_batch_to_connector=lambda _batch: phases.append("report"), + ) + }, + maybe_fit_token_budget=lambda _batch: phases.append("trim"), + ) + ResourceManager.prepare_resources(manager, batch) + assert phases == ["prepare", "trim", "report"] + assert context.context_chunk_size == 128 + assert all(r.py_verify_len == 3 for r in batch.generation_requests) + + +def test_all_decode_keeps_original_policy_and_device_mode(monkeypatch): + group, votes = _group(monkeypatch, all_decode=True) + for mixed, batch, worker, planner, _ in group: + before = [r.__dict__.copy() for r in batch.generation_requests] + assert not mixed.apply(batch, worker, planner, votes, 10) + assert [r.__dict__ for r in batch.generation_requests] == before + assert planner.device_windows + + +def test_optional_suffix_preserves_legacy_payload(monkeypatch): + group, votes = _group(monkeypatch) + mixed = group[0][0] + legacy, suffix = mixed.split_votes([[1, 2, 3] + vote for vote in votes], 3) + assert legacy == [[1, 2, 3]] * 8 + assert suffix == votes + with pytest.raises(ValueError, match="shape differs"): + mixed.split_votes([[1, 2, 3] + vote[:-1] for vote in votes], 3) + + +def test_unsupported_backend_declines_before_reading_snapshot(monkeypatch): + group, _ = _group(monkeypatch) + mixed, batch, worker, planner, _ = group[1] + mixed.runner._config.prefill_cuda_graph_backend = SimpleNamespace(name="PIECEWISE") + planner._mixed_current_snapshot_meta = None + vote = mixed.local_vote(batch, worker, planner, 10, 1, executor_eligible=True) + assert len(vote) == runtime.WIRE_WORDS + assert vote[5] == 0 + + +def test_host_metadata_skips_device_prologue_and_local_full_graph(): + runner = SimpleNamespace( + _dspark_trims_submitted_tokens=True, + _dspark_device_windows=True, + cuda_graph_runner=SimpleNamespace(agreed_ragged_bucket=192), + get_runtime_tokens_per_gen_step=lambda k: k + 1, + _ragged_token_lens=DecoderRunner._ragged_token_lens, + ) + metadata = SimpleNamespace( + runtime_tokens_per_gen_step=None, ragged_verify_lens=None, device_windows_mode=False + ) + requests = [SimpleNamespace(py_verify_len=1), SimpleNamespace(py_verify_len=3)] + DecoderRunner._publish_gen_token_layout( + runner, metadata, requests, runtime_draft_len=3, host_authoritative=True + ) + assert not metadata.device_windows_mode + assert metadata.ragged_verify_lens == [2, 4] + runner._dspark_host_window_step, runner._dspark_device_budget = True, 5 + assert not DecoderRunner._apply_device_window_prologue(runner, {}, None) + assert runner._dspark_device_budget is None + batch = SimpleNamespace(can_run_cuda_graph=True) + graph = SimpleNamespace( + _dspark_host_window_batch=batch, + is_encoder_decoder=False, + enable_encoder_decoder_mixed_cuda_graph=False, + ) + assert not CUDAGraphRunner._can_run_cuda_graph_batch(graph, batch) + + +def test_producer_stamp_copy_uses_same_event_and_validates_before_rotation(monkeypatch): + monkeypatch.setattr(torch.cuda, "Event", _Event) + planner = DSparkVerifyPlanner(cfg=DSparkScheduleConfig(3), device_windows=True) + planner.stage_confidence(torch.ones(2, 3), torch.tensor([1, 2], dtype=torch.int32)) + first = (planner._host_buffer, planner._host_confidence_stamp, planner._copy_event) + with pytest.raises(ValueError, match="one per slot"): + planner.stage_confidence(torch.zeros(2, 3), torch.ones(3, dtype=torch.int32)) + assert planner._host_buffer is first[0] + assert planner._host_confidence_stamp is first[1] + planner.stage_confidence(torch.zeros(2, 3), torch.tensor([3, 4], dtype=torch.int32)) + assert planner._prev_buffer is first[0] + assert planner._prev_confidence_stamp is first[1] + assert planner._prev_event is first[2] + assert planner._host_confidence_stamp.tolist() == [3, 4] + + +def test_forward_attempt_alone_is_not_confidence_publication(monkeypatch): + group, _ = _group(monkeypatch) + _, batch, worker, planner, _ = group[1] + runtime.stage_mixed_ownership(planner, worker, 10, batch.generation_requests) + assert planner._mixed_current_snapshot_meta["producer_attempts"] == {} + worker._draft_seq_host += 1 + runtime.stage_mixed_ownership(planner, worker, 11, batch.generation_requests) + assert len(planner._mixed_current_snapshot_meta["producer_attempts"]) == 64 + assert ( + planner._mixed_current_snapshot_meta["confidence_stamp"] is planner._host_confidence_stamp + ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sampling_admission.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sampling_admission.py new file mode 100644 index 000000000000..962cf4e5a7b6 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sampling_admission.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unsupported confidence/sampler combinations fail before model execution.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm._torch.speculative.dspark import validate_confidence_request +from tensorrt_llm._torch.speculative.interface import SpecMetadata + +pytestmark = pytest.mark.cpu_only + + +@pytest.mark.parametrize("enabled,rejection", [(False, True), (True, False), (True, True)]) +@pytest.mark.parametrize( + "controls,greedy", + [ + ({}, True), + ({"temperature": 0.0}, True), + ({"temperature": 1.0}, False), + ({"temperature": 1.0, "top_k": 1}, True), + ({"min_p": 0.2}, False), + ({"min_p": 1.0}, True), + ], +) +def test_request_admission_matches_actual_draft_sampling_controls( + enabled, rejection, controls, greedy +): + config = SimpleNamespace(enable_confidence_scheduling=enabled, use_rejection_sampling=rejection) + sampling = SimpleNamespace(temperature=None, top_k=None, top_p=None, min_p=None) + sampling.__dict__.update(controls) + request = SimpleNamespace(sampling_config=sampling, py_logits_post_processors=None) + if enabled and rejection and not greedy: + with pytest.raises(ValueError, match="advanced rejection draft sampling"): + validate_confidence_request(config, request) + else: + validate_confidence_request(config, request) + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_per_request_logit_processor_is_rejected_at_confidence_admission(enabled): + config = SimpleNamespace(enable_confidence_scheduling=enabled, use_rejection_sampling=False) + request = SimpleNamespace(sampling_config=None, py_logits_post_processors=[object()]) + if enabled: + with pytest.raises(ValueError, match="per-request logits post-processors"): + validate_confidence_request(config, request) + else: + validate_confidence_request(config, request) + + +@pytest.mark.parametrize( + "enabled,rejection,flags", + [ + (True, True, [True, True]), + (True, True, [True, False]), + (True, True, [False, True]), + (True, True, [False]), + (True, False, [False, True]), + (False, True, [False, True]), + ], +) +def test_host_sampler_guard_is_group_uniform_before_graph_lookup(enabled, rejection, flags): + errors = [] + for local_flag in flags: + runner = object.__new__(DecoderRunner) + runner._dspark_confidence_enabled = enabled + runner.mapping = SimpleNamespace(enable_lm_head_tp_in_adp=False, tp_size=len(flags)) + exchange = Mock(return_value=torch.tensor([[flag] for flag in flags], dtype=torch.int64)) + runner.dist = SimpleNamespace(tp_allgather_int64=exchange) + metadata = object.__new__(SpecMetadata) + metadata.use_rejection_sampling = rejection + metadata.is_all_greedy_sample = local_flag + metadata.group_all_greedy_sample = None + runner._sync_group_all_greedy_sample(metadata) + if enabled and rejection and not all(flags): + with pytest.raises(ValueError, match="advanced rejection draft sampling"): + runner._validate_dspark_draft_sampling(metadata) + errors.append(True) + else: + runner._validate_dspark_draft_sampling(metadata) + errors.append(False) + assert exchange.call_count == int(enabled and rejection and len(flags) > 1) + assert len(set(errors)) == 1 + + +@pytest.mark.parametrize( + "enabled,rejection", [(False, False), (False, True), (True, False), (True, True)] +) +def test_capture_gate_retains_only_supported_advanced_variants(enabled, rejection): + runner = object.__new__(DecoderRunner) + runner._dspark_confidence_enabled = enabled + runner._config = SimpleNamespace(spec_config=SimpleNamespace(use_rejection_sampling=rejection)) + assert runner._supports_advanced_draft_sampling() is not (enabled and rejection) + + +@pytest.mark.parametrize("processors", [False, True]) +def test_invalid_confidence_request_is_excluded_before_activation(processors): + executor = object.__new__(PyExecutor) + executor.model_engine = SimpleNamespace( + _dspark_confidence_enabled=True, + spec_config=SimpleNamespace(enable_confidence_scheduling=True, use_rejection_sampling=True), + ) + request = SimpleNamespace( + sampling_config=SimpleNamespace(temperature=1.0, top_k=None, top_p=None, min_p=None), + py_logits_post_processors=[object()] if processors else None, + ) + already_active = object() + executor.waiting_queue = object() + executor.active_requests = [already_active] + executor._fetch_new_requests = Mock(return_value=[request]) + executor._handle_errors = Mock() + executor._configure_csa2_remote_tail = Mock() + + assert executor._fetch_and_activate_new_requests() == [] + assert executor.active_requests == [already_active] + executor._configure_csa2_remote_tail.assert_not_called() + executor._handle_errors.assert_called_once() + assert executor._handle_errors.call_args.kwargs == { + "requests": [request], + "charge_budget": False, + } + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_feature_off_retains_existing_lm_head_tp_group_sync(enabled): + runner = object.__new__(DecoderRunner) + runner._dspark_confidence_enabled = enabled + runner.mapping = SimpleNamespace(enable_lm_head_tp_in_adp=True, tp_size=2) + exchange = Mock(return_value=torch.tensor([[True], [False]], dtype=torch.int64)) + runner.dist = SimpleNamespace(tp_allgather_int64=exchange) + metadata = object.__new__(SpecMetadata) + metadata.use_rejection_sampling = True + metadata.is_all_greedy_sample = True + metadata.group_all_greedy_sample = None + + runner._sync_group_all_greedy_sample(metadata) + + exchange.assert_called_once_with([True]) + assert not metadata.group_all_greedy_sample + assert not metadata.is_all_greedy_sample diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_schedule_fused.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_schedule_fused.py new file mode 100644 index 000000000000..7cbe2eb1711f --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_schedule_fused.py @@ -0,0 +1,543 @@ +# 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 + +import tensorrt_llm._torch.speculative.dspark_select as device_select +import tensorrt_llm._torch.speculative.dspark_schedule as dspark_schedule +from tensorrt_llm._torch.speculative.dspark_select import select_windows_device +from tensorrt_llm._torch.speculative.dspark_schedule import ( + DSparkScheduleConfig, + schedule_verify_lens_topk, + schedule_verify_lens_topk_fused_fill, +) +from tensorrt_llm._torch.speculative.ragged_helpers import fill_bucket_device + + +def _legacy_schedule_and_fill( + *, + survival: torch.Tensor, + budget: int, + num_real: int, + pad_len: int, + graph_num_tokens: int, + cfg: DSparkScheduleConfig, +) -> torch.Tensor: + rows = torch.arange(survival.shape[0], device=survival.device) + real_survival = torch.where( + (rows < num_real).unsqueeze(1), + survival, + torch.zeros_like(survival), + ) + scheduled = schedule_verify_lens_topk(survival=real_survival, budget=budget, cfg=cfg) + return fill_bucket_device( + scheduled + 1, + num_real=torch.tensor(num_real, device=survival.device), + graph_num_tokens=graph_num_tokens, + max_verify_len=cfg.resolved_max_verify_len + 1, + pad_fill=pad_len, + ) + + +@pytest.mark.parametrize( + "num_rows,num_real,budget,survival_eps", + [ + (1, 0, 0, 1e-6), + (1, 1, 0, 1e-6), + (2, 1, 1, 1e-6), + (2, 1, 9, 1e-6), + (4, 2, 3, 1e-6), + (4, 4, 16, 1e-6), + (8, 4, 9, 1e-6), + (2, 2, 1, 0.35), + (4, 1, 3, 0.35), + (8, 8, 32, 0.35), + ], +) +def test_cpu_fused_helper_matches_established_schedule_and_fill( + num_rows, num_real, budget, survival_eps +): + cfg = DSparkScheduleConfig( + block_size=5, + min_verify_len=1, + max_verify_len=5, + survival_eps=survival_eps, + ) + generator = torch.Generator().manual_seed( + 10000 * num_rows + 1000 * num_real + 10 * budget + int(survival_eps * 10) + ) + survival = torch.cumprod(torch.rand(num_rows, 5, generator=generator), dim=1) + survival[num_real:] = 0 + pad_len = 1 + token_floor = cfg.min_verify_len + 1 + max_token_len = cfg.resolved_max_verify_len + 1 + pad_tokens = (num_rows - num_real) * pad_len + capacity = min(max(budget, 0), num_real * cfg.schedulable_per_request) + minimum = num_real * token_floor + pad_tokens + maximum = num_real * max_token_len + pad_tokens + graph_num_tokens = min(maximum, minimum + capacity + min(num_real, 2)) + + expected = _legacy_schedule_and_fill( + survival=survival, + budget=budget, + num_real=num_real, + pad_len=pad_len, + graph_num_tokens=graph_num_tokens, + cfg=cfg, + ) + actual = schedule_verify_lens_topk_fused_fill( + survival=survival, + budget=budget, + num_real=num_real, + pad_len=pad_len, + graph_num_tokens=graph_num_tokens, + cfg=cfg, + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("num_real,budget,graph_tokens", [(2, 2, 9), (4, 4, 14)]) +def test_fused_workspace_selection_matches_tensor_and_reuses_storage( + num_real, budget, graph_tokens +): + cfg = DSparkScheduleConfig(block_size=3, min_verify_len=1) + workspace = device_select.DeviceWindowWorkspace.allocate( + max_rows=6, max_tokens=16, device="cpu" + ) + workspace.verify_lens.fill_(-777) + pointers = { + name: getattr(workspace, name).data_ptr() + for name in ("verify_lens", "qo_indptr", "req_idx", "kv_correction") + } + logits = torch.tensor([[10.0, 5.0, 1.0], [-2.0, -3.0, -4.0], [0.5, 0.2, 0.1], [1.0, 2.0, 3.0]]) + for confidence in (logits, logits.flip(0)): + kwargs = dict( + confidence_logits=confidence, + slot_idx=torch.tensor([0, 1, 2, 3]), + num_real=num_real, + budget=budget, + graph_num_tokens=graph_tokens, + cfg=cfg, + pad_len=1, + ) + expected = select_windows_device(**kwargs, use_fused_exact=False) + actual = select_windows_device(**kwargs, use_fused_exact=True, workspace=workspace) + assert expected.workspace is None + assert actual.workspace is workspace + for name, pointer in pointers.items(): + result = getattr(actual, name) + assert result.data_ptr() == pointer + torch.testing.assert_close(result, getattr(expected, name), rtol=0, atol=0) + assert expected.verify_lens.data_ptr() != workspace.verify_lens.data_ptr() + assert workspace.verify_lens[4:].tolist() == [-777, -777] + + +@pytest.mark.parametrize("num_rows", [0, 2]) +def test_fused_helper_preserves_output_address_and_unused_capacity(num_rows): + cfg = DSparkScheduleConfig(block_size=3, min_verify_len=1) + output = torch.full((num_rows + 2,), -777, dtype=torch.int32) + actual = schedule_verify_lens_topk_fused_fill( + survival=torch.ones(num_rows, 3), + budget=0, + num_real=num_rows, + pad_len=1, + graph_num_tokens=2 * num_rows, + cfg=cfg, + out=output, + ) + assert actual.untyped_storage().data_ptr() == output.untyped_storage().data_ptr() + assert actual.tolist() == [2] * num_rows + assert output[num_rows:].tolist() == [-777, -777] + + +@pytest.mark.parametrize( + "output,error,message", + [ + (torch.empty(2, dtype=torch.int32, device="meta"), ValueError, "share survival.device"), + (torch.empty(2, dtype=torch.int64), TypeError, "dtype torch.int32"), + (torch.empty(1, 2, dtype=torch.int32), ValueError, "1-D tensor"), + (torch.empty(1, dtype=torch.int32), ValueError, "at least 2 elements"), + (torch.empty(4, dtype=torch.int32)[::2], ValueError, "contiguous"), + ], +) +def test_fused_helper_rejects_incompatible_output_before_dispatch(output, error, message): + cfg = DSparkScheduleConfig(block_size=3, min_verify_len=1) + with pytest.raises(error, match=message): + schedule_verify_lens_topk_fused_fill( + survival=torch.ones(2, 3), + budget=0, + num_real=2, + pad_len=1, + graph_num_tokens=4, + cfg=cfg, + out=output, + ) + + +def test_round_robin_slack_is_not_promoted_into_confidence_budget(): + cfg = DSparkScheduleConfig( + block_size=5, + min_verify_len=1, + max_verify_len=5, + survival_eps=0.5, + ) + survival = torch.tensor( + [ + [1.0, 0.9, 0.1, 0.1, 0.1], + [1.0, 0.8, 0.49, 0.1, 0.1], + ] + ) + actual = schedule_verify_lens_topk_fused_fill( + survival=survival, + budget=2, + num_real=2, + pad_len=1, + graph_num_tokens=7, + cfg=cfg, + ) + # Budget two selects one position per row, then the already-paid graph + # remainder goes to row zero. Promoting that remainder into top-k would + # instead select row one's 0.49 candidate and produce [3, 4]. + assert actual.tolist() == [4, 3] + + +def test_zero_epsilon_declines_fusion_and_preserves_tensor_semantics(monkeypatch): + cfg = DSparkScheduleConfig( + block_size=5, + min_verify_len=1, + max_verify_len=5, + survival_eps=0.0, + ) + + def _must_not_run(**_kwargs): + raise AssertionError("zero-epsilon selector must use the tensor path") + + monkeypatch.setattr(device_select, "schedule_verify_lens_topk_fused_fill", _must_not_run) + result = select_windows_device( + confidence_logits=torch.zeros(3, 5), + slot_idx=torch.tensor([0, 1, 2]), + num_real=2, + budget=2, + graph_num_tokens=7, + cfg=cfg, + pad_len=1, + use_fused_exact=True, + ) + assert int(result.verify_lens.sum()) == 7 + + +def test_tensor_controls_decline_fusion(monkeypatch): + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1, max_verify_len=5) + + def _must_not_run(**_kwargs): + raise AssertionError("tensor controls must use the capture-safe tensor path") + + monkeypatch.setattr(device_select, "schedule_verify_lens_topk_fused_fill", _must_not_run) + result = select_windows_device( + confidence_logits=torch.zeros(2, 5), + slot_idx=torch.tensor([0, 1]), + num_real=torch.tensor(2, dtype=torch.int64), + budget=torch.tensor(1, dtype=torch.int64), + graph_num_tokens=5, + cfg=cfg, + pad_len=1, + use_fused_exact=True, + ) + assert int(result.verify_lens.sum()) == 5 + + +def test_fused_helper_rejects_infeasible_graph_without_dispatch(): + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1, max_verify_len=5) + with pytest.raises(ValueError, match="cannot realize"): + schedule_verify_lens_topk_fused_fill( + survival=torch.ones(2, 5), + budget=4, + num_real=2, + pad_len=1, + graph_num_tokens=5, + cfg=cfg, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("graph_bs", [16, 32, 64, 128, 192, 256]) +def test_cuda_fused_matches_tensor_oracle_across_all_graph_sizes(graph_bs): + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1, max_verify_len=5) + generator = torch.Generator(device="cuda").manual_seed(347805 + graph_bs) + confidence = torch.rand((graph_bs, 5), generator=generator, device="cuda") + survival = torch.cumprod(confidence, dim=1) + # Include exact ties and non-finite candidates in every graph shape. + survival[:2] = 0.75 + survival[2, 2] = torch.nan + survival[3, 1] = torch.inf + num_real = graph_bs - 3 + survival[num_real:] = 0 + budget = min(2 * num_real, num_real * cfg.schedulable_per_request) + pad_len = 1 + minimum = num_real * (cfg.min_verify_len + 1) + (graph_bs - num_real) + graph_num_tokens = minimum + budget + num_real + + expected = _legacy_schedule_and_fill( + survival=survival, + budget=budget, + num_real=num_real, + pad_len=pad_len, + graph_num_tokens=graph_num_tokens, + cfg=cfg, + ) + actual = schedule_verify_lens_topk_fused_fill( + survival=survival, + budget=budget, + num_real=num_real, + pad_len=pad_len, + graph_num_tokens=graph_num_tokens, + cfg=cfg, + ) + torch.cuda.synchronize() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_model_engine_prewarm_requires_device_windows(): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + + engine = object.__new__(DecoderRunner) + engine._dspark_fused_scheduler_enabled = True + engine._dspark_confidence_enabled = True + engine._dspark_trims_submitted_tokens = True + engine._dspark_device_windows = False + + def _must_not_get_worker(): + raise AssertionError("disabled device-window mode must not prewarm fusion") + + engine._get_spec_worker = _must_not_get_worker + DecoderRunner._warmup_dspark_fused_scheduler(engine) + assert engine._dspark_fused_schedule_ready_sizes == set() + + +def test_model_engine_prewarm_records_only_successful_graph_sizes(monkeypatch): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import runner as decoder_runner + + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1, max_verify_len=5) + planner = SimpleNamespace( + cfg=cfg, + exact_cost_table=SimpleNamespace(tables={16: object(), 32: object()}), + ) + engine = object.__new__(DecoderRunner) + engine._dspark_fused_scheduler_enabled = True + engine._dspark_confidence_enabled = True + engine._dspark_trims_submitted_tokens = True + engine._dspark_device_windows = True + engine._get_spec_worker = lambda: SimpleNamespace(verify_planner=planner) + + original_ones = torch.ones + monkeypatch.setattr( + decoder_runner.torch, + "ones", + lambda shape, dtype=None, device=None: original_ones(shape, dtype=dtype), + ) + monkeypatch.setattr(decoder_runner.torch.cuda, "synchronize", lambda: None) + calls = [] + + def _fake_fused(**kwargs): + calls.append(kwargs["num_real"]) + if kwargs["num_real"] == 32: + raise RuntimeError("synthetic compile failure") + return torch.empty(kwargs["num_real"], dtype=torch.int32) + + monkeypatch.setattr(dspark_schedule, "schedule_verify_lens_topk_fused_fill", _fake_fused) + DecoderRunner._warmup_dspark_fused_scheduler(engine) + assert calls == [16, 32] + assert engine._dspark_fused_schedule_ready_sizes == {16} + + +def test_runtime_failure_retires_only_one_shape_and_retries_once(): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + from tensorrt_llm._torch.speculative.dspark_schedule import DSparkFusedScheduleError + + engine = object.__new__(DecoderRunner) + engine._dspark_fused_schedule_ready_sizes = {16, 32} + calls = [] + + def _select(**kwargs): + calls.append(kwargs["use_fused_exact"]) + if kwargs["use_fused_exact"]: + raise DSparkFusedScheduleError("synthetic launch failure") + return "tensor-result" + + result = DecoderRunner._select_dspark_windows_with_fused_fallback( + engine, + select_fn=_select, + selector_kwargs={}, + padded_bs=16, + ) + assert result == "tensor-result" + assert calls == [True, False] + assert engine._dspark_fused_schedule_ready_sizes == {32} + assert engine._dspark_fused_schedule_failure_counts == {16: 1} + + calls.clear() + result = DecoderRunner._select_dspark_windows_with_fused_fallback( + engine, + select_fn=_select, + selector_kwargs={}, + padded_bs=16, + ) + assert result == "tensor-result" + assert calls == [False] + assert engine._dspark_fused_schedule_failure_counts == {16: 1} + + +def _tensor_prologue_case(): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + + confidence = torch.tensor([[10.0, 10.0, 10.0], [-10.0, -10.0, -10.0]]) + planner = SimpleNamespace( + cfg=DSparkScheduleConfig(block_size=3, min_verify_len=1), + apply_calibration=None, + ) + worker = SimpleNamespace( + verify_planner=planner, + staged_confidence_buffer=lambda: confidence, + batch_slot_view=lambda rows: torch.arange(rows), + verified_draft_seq_cuda=lambda: None, + ) + engine = object.__new__(DecoderRunner) + engine.cuda_graph_runner = SimpleNamespace(agreed_ragged_bucket=6, ragged_pad_verify_len=0) + engine._config = SimpleNamespace(use_mrope=False) + engine._dspark_device_budget = 2 + engine._dspark_prev_covers_batch = True + engine._dspark_last_padded_bs = 2 + engine._dspark_last_num_real = 2 + engine._dspark_fused_schedule_ready_sizes = set() + engine._get_spec_worker = lambda: worker + engine.ragged_verify_lens_cuda = torch.tensor([3, 3], dtype=torch.int32) + engine.ragged_qo_indptr_cuda = torch.tensor([0, 3, 6], dtype=torch.int32) + engine.previous_batch_indices_cuda = torch.tensor([0, 1], dtype=torch.int64) + engine.input_ids_cuda = torch.full((6,), -1, dtype=torch.int32) + engine.position_ids_cuda = torch.tensor([10, 11, 12, 20, 21, 22], dtype=torch.int32) + engine.previous_pos_indices_cuda = torch.full((6,), -1, dtype=torch.int64) + engine.previous_pos_id_offsets_cuda = torch.full((6,), -1, dtype=torch.int32) + engine.previous_kv_lens_offsets_cuda = torch.tensor([-1, -2], dtype=torch.int32) + engine.draft_tokens_cuda = torch.full((4,), -1, dtype=torch.int32) + events = [] + attn_metadata = SimpleNamespace( + num_contexts=0, + kv_lens_cuda=torch.tensor([16, 26], dtype=torch.int32), + apply_device_ragged_layout=lambda lens, owners, correction: events.append( + ("attention", lens.clone(), owners.clone(), correction.clone()) + ), + ) + spec_metadata = SimpleNamespace( + remap_expanded_sampling_params=lambda owners, count: events.append( + ("sampling", owners.clone(), count) + ), + ) + new_tensors = SimpleNamespace( + new_tokens=torch.tensor([[100, 200], [101, 201], [102, 202], [103, 203]]), + new_tokens_lens=torch.tensor([2, 1], dtype=torch.int32), + next_draft_tokens=torch.tensor([[11, 12, 13], [21, 22, 23]]), + ) + return ( + engine, + {"spec_metadata": spec_metadata, "attn_metadata": attn_metadata}, + new_tensors, + events, + ) + + +def test_tensor_prologue_rebuilds_owners_positions_and_both_kv_terms(): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + + engine, inputs, new_tensors, events = _tensor_prologue_case() + lens_address = engine.ragged_verify_lens_cuda.data_ptr() + qo_address = engine.ragged_qo_indptr_cuda.data_ptr() + assert DecoderRunner._apply_device_window_prologue(engine, inputs, new_tensors) + assert engine._dspark_device_budget is None + assert engine.ragged_verify_lens_cuda.tolist() == [4, 2] + assert engine.ragged_qo_indptr_cuda.tolist() == [0, 4, 6] + assert engine.ragged_verify_lens_cuda.data_ptr() == lens_address + assert engine.ragged_qo_indptr_cuda.data_ptr() == qo_address + assert engine.input_ids_cuda.tolist() == [100, 101, 102, 103, 200, 201] + assert engine.position_ids_cuda.tolist() == [10, 11, 12, 13, 20, 21] + assert engine.previous_pos_indices_cuda.tolist() == [0, 0, 0, 0, 1, 1] + assert engine.previous_pos_id_offsets_cuda.tolist() == [2, 2, 2, 2, 1, 1] + assert engine.draft_tokens_cuda.tolist() == [11, 12, 13, 21] + assert inputs["attn_metadata"].kv_lens_cuda.tolist() == [18, 24] + assert engine.previous_kv_lens_offsets_cuda.tolist() == [-2, -1] + assert ( + inputs["attn_metadata"].kv_lens_cuda + engine.previous_kv_lens_offsets_cuda + ).tolist() == [16, 23] + assert [event[0] for event in events] == ["sampling", "attention"] + assert events[0][1].tolist() == [0, 0, 0, 0, 1, 1] + assert events[0][2] == 6 + + +def test_decode_only_prologue_declines_local_prefill_without_mutating_buffers(): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + + engine, inputs, new_tensors, events = _tensor_prologue_case() + inputs["attn_metadata"].num_contexts = 1 + tensors = ( + engine.input_ids_cuda, + engine.position_ids_cuda, + engine.ragged_verify_lens_cuda, + engine.ragged_qo_indptr_cuda, + inputs["attn_metadata"].kv_lens_cuda, + engine.previous_kv_lens_offsets_cuda, + ) + originals = [tensor.clone() for tensor in tensors] + assert not DecoderRunner._apply_device_window_prologue(engine, inputs, new_tensors) + for tensor, original in zip(tensors, originals): + torch.testing.assert_close(tensor, original, rtol=0, atol=0) + assert not events + + +@pytest.mark.parametrize("confidence,fused", [(False, False), (True, False), (True, True)]) +def test_workspace_allocation_follows_decoder_state_initialization(monkeypatch, confidence, fused): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import runner as decoder_runner + + original_empty, original_zeros = torch.empty, torch.zeros + + def cpu_empty(*args, **kwargs): + kwargs["device"] = "cpu" + return original_empty(*args, **kwargs) + + def cpu_zeros(*args, **kwargs): + kwargs["device"] = "cpu" + return original_zeros(*args, **kwargs) + + monkeypatch.setattr(torch, "empty", cpu_empty) + monkeypatch.setattr(torch, "zeros", cpu_zeros) + monkeypatch.setattr(decoder_runner, "prefer_pinned", lambda: False) + monkeypatch.setattr(decoder_runner, "LoraParamBuilder", lambda **kwargs: SimpleNamespace()) + engine = object.__new__(DecoderRunner) + engine._config = SimpleNamespace( + is_spec_decode=True, + max_draft_loop_tokens=3, + max_batch_size=4, + max_num_tokens=16, + max_beam_width=1, + max_seq_len=16, + use_mrope=False, + attention_backend="TRTLLM", + spec_config=SimpleNamespace( + enable_confidence_scheduling=confidence, + enable_fused_confidence_scheduler=fused, + ), + ) + engine._compute_dynamic_draft_len_mapping = lambda: {} + # Follow the constructor order: buffers exist before mutable state/flags. + engine._allocate_decoder_buffers() + engine._init_decoder_state() + if confidence and fused: + workspace = engine._dspark_device_window_workspace + assert workspace.verify_lens.numel() == 5 + assert workspace.req_idx.numel() == 16 + assert workspace.verify_lens.device == engine.input_ids_cuda.device + else: + assert engine._dspark_device_window_workspace is None diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_select.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_select.py index 0d1d9d737417..aa371b548346 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_select.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_select.py @@ -7,11 +7,41 @@ from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig from tensorrt_llm._torch.speculative.dspark_select import ( + DeviceWindowWorkspace, gather_packed_draft_tokens, select_windows_device, ) +def test_workspace_slices_keep_persistent_storage(): + workspace = DeviceWindowWorkspace.allocate(max_rows=4, max_tokens=16, device="cpu") + owners = [ + workspace.verify_lens, + workspace.qo_indptr, + workspace.req_idx, + workspace.kv_correction, + workspace.past_seen, + ] + addresses = [tensor.data_ptr() for tensor in owners] + workspace.validate(device=torch.device("cpu"), num_rows=2, num_tokens=8) + workspace.validate(device=torch.device("cpu"), num_rows=4, num_tokens=16) + assert [tensor.data_ptr() for tensor in owners] == addresses + + +@pytest.mark.parametrize("num_rows,num_tokens", [(5, 16), (4, 17)]) +def test_workspace_rejects_out_of_capacity_shape(num_rows, num_tokens): + workspace = DeviceWindowWorkspace.allocate(max_rows=4, max_tokens=16, device="cpu") + with pytest.raises(ValueError, match="capacity"): + workspace.validate(device=torch.device("cpu"), num_rows=num_rows, num_tokens=num_tokens) + + +def test_workspace_rejects_wrong_row_extent_dtype(): + workspace = DeviceWindowWorkspace.allocate(max_rows=4, max_tokens=16, device="cpu") + workspace.verify_lens = workspace.verify_lens.to(torch.int64) + with pytest.raises(TypeError, match="verify_lens must have dtype"): + workspace.validate(device=torch.device("cpu"), num_rows=4, num_tokens=16) + + def test_fresh_confidence_spends_budget_on_the_best_real_row(): result = select_windows_device( confidence_logits=torch.tensor( diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts.py index 04ce5bc4a276..0202110e8d7f 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts.py @@ -11,6 +11,8 @@ resolve_confidence_head, ) +pytestmark = pytest.mark.cpu_only + @pytest.mark.parametrize("key", ["sts_temperatures", "temperatures"]) def test_loads_native_and_sglang_temperature_keys(tmp_path, key): @@ -51,6 +53,14 @@ def test_rejects_non_list_temperatures(tmp_path, key, temperatures): load_sts_temperatures_from_path(str(path)) +@pytest.mark.parametrize("temperature", [0.0, -1.0, float("nan"), float("inf"), -float("inf")]) +def test_rejects_nonfinite_or_nonpositive_temperatures(tmp_path, temperature): + path = tmp_path / "sts.json" + path.write_text(json.dumps({"temperatures": [1.0, temperature]}), encoding="utf-8") + with pytest.raises(ValueError, match="finite positive"): + load_sts_temperatures_from_path(str(path)) + + class _Head: pass 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..fdfcca0155b2 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -137,7 +137,11 @@ def test_worker_lazy_init_window_buffers(): 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 + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import cuda_graph_dummy_request_id + + assert worker._graph_dummy_id_floor == cuda_graph_dummy_request_id( + worker.max_draft_len, variant=1, max_draft_len=worker.max_draft_len + ) # 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 @@ -422,17 +426,22 @@ def test_prepare_assigns_slots_to_disagg_generation_requests(): assert worker._batch_to_slot[:2].tolist() == [s0, s1] -def test_prepare_keeps_dummy_generation_requests_on_scratch_row(): +@pytest.mark.parametrize("physical_k", [3, 4, 5]) +@pytest.mark.parametrize("variant", [0, 1]) +def test_prepare_keeps_dummy_generation_requests_on_scratch_row(physical_k, variant): """ADP-idle (id 0) and CUDA-graph padding dummies never consume a real slot.""" - from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import cuda_graph_dummy_request_id from tensorrt_llm._torch.pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID worker = _make_worker() + worker.spec_config.max_draft_len = physical_k meta = _make_metadata(max_num_requests=4) - worker._lazy_init(_fake_draft_model(), meta) + draft_model = _fake_draft_model() + draft_model.block_size = physical_k + worker._lazy_init(draft_model, meta) meta._dspark_worker = worker - graph_dummy = CUDA_GRAPH_DUMMY_REQUEST_ID - worker.max_draft_len + graph_dummy = cuda_graph_dummy_request_id(physical_k, variant=variant, max_draft_len=physical_k) meta.request_ids = [1000, ATTENTION_DP_DUMMY_REQUEST_ID, graph_dummy] meta.num_generations = 3 meta.prepare() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_stamp_wrap_cpu.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_stamp_wrap_cpu.py new file mode 100644 index 000000000000..d7101c14894f --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_stamp_wrap_cpu.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Draft-stamp scalar bounds and neutral reset without device allocation.""" + +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark import DSv4DSparkWorker +from tensorrt_llm._torch.speculative.dspark_schedule import NEUTRAL_CONFIDENCE_LOGIT + +pytestmark = pytest.mark.cpu_only + + +@pytest.mark.parametrize("previous", [0, (1 << 31) - 2, (1 << 31) - 1]) +def test_stamp_sequence_stays_representable_and_wraps_neutrally(previous): + worker = object.__new__(DSv4DSparkWorker) + worker._draft_seq_host = previous + worker._draft_seq_cuda = torch.tensor(previous, dtype=torch.int32) + worker._confidence_stamp = torch.tensor([previous, 0], dtype=torch.int32) + worker._confidence_logits = torch.full((2, 3), -3.0) + pointers = [ + worker._draft_seq_cuda.data_ptr(), + worker._confidence_stamp.data_ptr(), + worker._confidence_logits.data_ptr(), + ] + + assert worker.bump_draft_seq() == previous + current = (previous + 1) & ((1 << 31) - 1) + assert worker._draft_seq_host == worker._draft_seq_cuda.item() == current + assert worker.verified_draft_seq_cuda().item() == current - 1 + if current == 0: + assert worker._confidence_stamp.tolist() == [0, 0] + assert torch.all(worker._confidence_logits == NEUTRAL_CONFIDENCE_LOGIT) + # A real draft emitted at sequence0 is fresh on the next step. Rows + # without that draft remain neutral even if their reset stamp matches. + worker._confidence_logits[0].fill_(-2.0) + assert worker.bump_draft_seq() == 0 + assert worker.verified_draft_seq_cuda().item() == 0 + assert torch.all(worker._confidence_logits[0] == -2.0) + assert torch.all(worker._confidence_logits[1] == NEUTRAL_CONFIDENCE_LOGIT) + else: + assert worker._confidence_stamp.tolist() == [previous, 0] + assert torch.all(worker._confidence_logits == -3.0) + assert pointers == [ + worker._draft_seq_cuda.data_ptr(), + worker._confidence_stamp.data_ptr(), + worker._confidence_logits.data_ptr(), + ] + + +def test_stamp_wrap_without_allocated_confidence_buffers(): + worker = object.__new__(DSv4DSparkWorker) + worker._draft_seq_host = (1 << 31) - 1 + worker._draft_seq_cuda = None + worker._confidence_stamp = None + worker._confidence_logits = None + assert worker.bump_draft_seq() == (1 << 31) - 1 + assert worker._draft_seq_host == 0 + assert worker.verified_draft_seq_cuda() is None diff --git a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py index 0c9468c48a75..6bb25b0734e3 100644 --- a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py +++ b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py @@ -16,6 +16,7 @@ from tensorrt_llm._utils import get_sm_version _ROPE_DIM = 64 +_PHYSICAL_KS = (1, 2, 3, 4, 5, 6, 7, 8) pytestmark = pytest.mark.skipif( not torch.cuda.is_available() @@ -106,7 +107,7 @@ def _reference( ("invalid_case", "reason"), [ ("q_dtype", "q dtype must be BF16"), - ("block_size", "draft block size must be 5 or 6"), + ("block_size", "draft block size must be 1 through 8"), ("valid_len_dtype", "valid_len must use INT64"), ("freqs_shape", "inverse_rope_freqs shape must be"), ], @@ -119,8 +120,8 @@ def test_fused_dsv4_dspark_attention_support_gate_logs_rejection(monkeypatch, in if invalid_case == "q_dtype": inputs[0] = q.float() elif invalid_case == "block_size": - inputs[0] = q[:, :4].contiguous() - inputs[4] = freqs[:, :4].contiguous() + inputs[0] = q.new_empty((q.shape[0], 9, *q.shape[2:])) + inputs[4] = freqs.new_empty((q.shape[0], 9, *freqs.shape[2:])) elif invalid_case == "valid_len_dtype": inputs[2] = inputs[2].to(torch.int32) else: @@ -177,7 +178,7 @@ def test_fused_dsv4_dspark_attention_rejects_invalid_inputs_before_launch( ) -@pytest.mark.parametrize("block", (5, 6)) +@pytest.mark.parametrize("block", _PHYSICAL_KS) @pytest.mark.parametrize( ("start_pos_values", "valid_len_values"), [ @@ -188,8 +189,10 @@ def test_fused_dsv4_dspark_attention_rejects_invalid_inputs_before_launch( ([5, 100], [6, 101]), # Bootstrapped positions with short physical suffixes exercise wraparound. ([257, 390], [3, 5]), + # The first row is a graph-padding request with no initialized context. + ([0, 5], [0, 6]), ], - ids=["full_window", "partial_window", "wrapped_suffix"], + ids=["full_window", "partial_window", "wrapped_suffix", "empty_scratch_window"], ) def test_fused_dsv4_dspark_attention_matches_reference(block, start_pos_values, valid_len_values): from tensorrt_llm._torch.custom_ops.dspark_attention_custom_op import ( @@ -199,6 +202,9 @@ def test_fused_dsv4_dspark_attention_matches_reference(block, start_pos_values, q, main_kv, block_kv, kv_cache, slots, start_pos, sink, blk_freqs, inverse_rope_freqs = ( _make_inputs(29, block=block, start_pos_values=start_pos_values) ) + if valid_len_values[0] == 0: + # Use a reserved scratch page rather than aliasing a live request slot. + slots[0] = kv_cache.shape[0] - 1 valid_len = torch.tensor(valid_len_values, device=q.device, dtype=torch.long) expected, expected_cache = _reference( q, @@ -215,6 +221,8 @@ def test_fused_dsv4_dspark_attention_matches_reference(block, start_pos_values, draft_block, slots_i32, cache_seqs = _prepare_attention_inputs( main_kv, block_kv, kv_cache, slots, start_pos ) + # A wrong draft mask must not accidentally pass because padding is zero. + draft_block[:, block:].fill_(17) actual = fused_dsv4_dspark_attention( q, @@ -232,7 +240,8 @@ def test_fused_dsv4_dspark_attention_matches_reference(block, start_pos_values, torch.testing.assert_close(kv_cache, expected_cache, rtol=0, atol=0) -def test_fused_dsv4_dspark_attention_cuda_graph_replay(): +@pytest.mark.parametrize("block", _PHYSICAL_KS) +def test_fused_dsv4_dspark_attention_cuda_graph_replay(block): from tensorrt_llm._torch.custom_ops.dspark_attention_custom_op import ( fused_dsv4_dspark_attention, warmup_fused_dsv4_dspark_attention, @@ -244,9 +253,9 @@ def test_fused_dsv4_dspark_attention_cuda_graph_replay(): ) q, main_kv, block_kv, kv_cache, slots, start_pos, sink, blk_freqs, inverse_rope_freqs = ( - _make_inputs(3, start_pos_values=[257, 5, 390]) + _make_inputs(3, block=block, start_pos_values=[257, 5, 390, 0]) ) - batch, block, dim = block_kv.shape + batch, _, dim = block_kv.shape scale = dim**-0.5 main_x = main_kv.unsqueeze(1) block_x = block_kv @@ -256,8 +265,10 @@ def test_fused_dsv4_dspark_attention_cuda_graph_replay(): main_freqs[..., 0] = 1 block_freqs[..., 0] = 1 slots = slots.long() + slots[-1] = kv_cache.shape[0] - 1 start_pos = start_pos.long() valid_len = _legacy_valid_len(start_pos) + valid_len[-1] = 0 # The named best-effort prewarm calls the same self-JIT ops used below, # so graph capture sees only hot compile-cache entries. @@ -282,7 +293,15 @@ def test_fused_dsv4_dspark_attention_cuda_graph_replay(): main_x.copy_(torch.randn_like(main_x)) block_x.copy_(torch.randn_like(block_x)) - valid_len.copy_(torch.tensor([3, 2, 7], device=q.device)) + q.copy_(torch.randn_like(q)) + sink.copy_(torch.randn_like(sink)) + # Replay must read new positions, page mappings, and inverse-RoPE phases. + start_pos.copy_(torch.tensor([390, 127, 257, 0], device=q.device)) + slots[:3].copy_(slots[:3].roll(1)) + valid_len.copy_(torch.tensor([3, 128, 7, 0], device=q.device)) + table = precompute_dspark_freqs_cis(_ROPE_DIM, 512, device=q.device) + blk_freqs = table[start_pos.unsqueeze(1) + 1 + torch.arange(block, device=q.device)] + inverse_rope_freqs.copy_(torch.view_as_real(blk_freqs)) expected_main = cute_dsl_dspark_rmsnorm_rope( main_x, weight, main_freqs, 1, _ROPE_DIM, 1e-6, True, True, False ).squeeze(1) @@ -312,13 +331,14 @@ def test_fused_dsv4_dspark_attention_cuda_graph_replay(): @pytest.mark.parametrize("persist", (False, True)) -def test_dspark_attention_forward_batched_matches_fallback(monkeypatch, persist): +@pytest.mark.parametrize("block", _PHYSICAL_KS) +def test_dspark_attention_forward_batched_matches_fallback(monkeypatch, persist, block): import tensorrt_llm._torch.models.modeling_dspark as dspark_attention torch.manual_seed(17) device = torch.device("cuda") dtype = torch.bfloat16 - batch, block, hidden = 2, 6, 64 + batch, hidden = 2, 64 heads, head_dim, rope_dim = 128, 512, 64 q_rank, groups, o_rank, window = 1024, 8, 32, 128 @@ -444,7 +464,8 @@ def test_warmup_is_best_effort(monkeypatch): assert "synthetic compile failure" in warnings[0] -def test_self_jit_reuses_one_kernel_across_runtime_shapes_and_scales(): +@pytest.mark.parametrize("block", _PHYSICAL_KS) +def test_self_jit_reuses_one_kernel_across_runtime_shapes_and_scales(block): from tensorrt_llm._torch.custom_ops.dspark_attention_custom_op import ( _dspark_attention_kernel_cache, _get_dspark_arch_str, @@ -459,13 +480,13 @@ def test_self_jit_reuses_one_kernel_across_runtime_shapes_and_scales(): assert _get_dspark_arch_str(101) is None assert _get_dspark_arch_str(109) is None assert _get_dspark_arch_str() in ("sm_100", "sm_103") - assert is_dsv4_dspark_attention_config_supported(5, 128, 512, 128) - assert not is_dsv4_dspark_attention_config_supported(4, 128, 512, 128) + assert is_dsv4_dspark_attention_config_supported(block, 128, 512, 128) + assert not is_dsv4_dspark_attention_config_supported(9, 128, 512, 128) assert not is_dsv4_dspark_attention_config_supported(5, 24, 512, 128) _dspark_attention_kernel_cache.clear() q, main_kv, block_kv, kv_cache, slots, start_pos, sink, _, freqs = _make_inputs( - 11, start_pos_values=[300, 4, 250], cache_pages=40 + 11, block=block, start_pos_values=[300, 4, 250], cache_pages=40 ) scale = q.shape[-1] ** -0.5 valid_len = _legacy_valid_len(start_pos) @@ -505,7 +526,9 @@ def test_self_jit_reuses_one_kernel_across_runtime_shapes_and_scales(): # hit on the same compiled object without compiler calls or runtime padding. for batch in (1, 3, 32): values = [5 + (37 * i) % 386 for i in range(batch)] - args = _make_inputs(17 + batch, start_pos_values=values, cache_pages=batch + 41) + args = _make_inputs( + 17 + batch, block=block, start_pos_values=values, cache_pages=batch + 41 + ) q_b, main_b, block_b, cache_b, slots_b, pos_b, sink_b, blk_freqs_b, freqs_b = args if batch == 3: cache_b = cache_b.clone() @@ -543,6 +566,43 @@ def test_self_jit_reuses_one_kernel_across_runtime_shapes_and_scales(): assert len(_dspark_attention_kernel_cache) == 1 +def test_attention_cache_isolates_all_k_with_shared_dummy_scratch(): + import tensorrt_llm._torch.custom_ops.dspark_attention_custom_op as attention_op + + attention_op._dspark_attention_kernel_cache.clear() + for block in _PHYSICAL_KS: + attention_op.warmup_fused_dsv4_dspark_attention(block, 1e-6) + arch = attention_op._get_dspark_arch_str() + kernels = attention_op._dspark_attention_kernel_cache.copy() + assert set(kernels) == {(block, arch) for block in _PHYSICAL_KS} + assert len({id(kernel) for kernel in kernels.values()}) == len(_PHYSICAL_KS) + + # Keep all specializations live while alternating K. Two padding requests + # read the reserved scratch page, never a real request's window. No duplicate + # cache-write launches are needed for this attention-only aliasing check. + for block in reversed(_PHYSICAL_KS): + q, main, block_kv, cache, slots, pos, sink, blk_freqs, freqs = _make_inputs( + 59, block=block, start_pos_values=[257, 0, 0] + ) + scratch = cache.shape[0] - 1 + slots.copy_(torch.tensor([0, scratch, scratch], device=q.device, dtype=slots.dtype)) + main[1:].zero_() + cache[0, 257 % 128].copy_(main[0]) + cache[scratch, 0].zero_() + valid = torch.tensor([3, 0, 0], device=q.device, dtype=torch.long) + draft = block_kv.new_zeros((3, 8, 512)) + draft[:, :block].copy_(block_kv) + expected, expected_cache = _reference( + q, main, block_kv, cache, slots, pos, valid, sink, blk_freqs, 512**-0.5 + ) + actual = attention_op.fused_dsv4_dspark_attention( + q, draft, cache, slots, pos, valid, sink, freqs, 512**-0.5 + ) + torch.testing.assert_close(actual, expected, rtol=5e-2, atol=3e-2) + torch.testing.assert_close(cache, expected_cache, rtol=0, atol=0) + assert attention_op._dspark_attention_kernel_cache == kernels + + def test_compile_without_real_specimens_and_cached_wrapper_avoids_views(monkeypatch): import tensorrt_llm._torch.custom_ops.dspark_attention_custom_op as dspark_attention_op @@ -596,11 +656,12 @@ def unexpected_host_op(*args, **kwargs): assert output.shape == q.shape -def test_attention_cache_miss_rejects_cuda_graph_capture(monkeypatch): +@pytest.mark.parametrize("block", _PHYSICAL_KS) +def test_attention_cache_miss_rejects_cuda_graph_capture(monkeypatch, block): import tensorrt_llm._torch.custom_ops.dspark_attention_custom_op as dspark_attention_op dspark_attention_op._dspark_attention_kernel_cache.clear() - q, main_kv, block_kv, kv_cache, slots, start_pos, sink, _, freqs = _make_inputs(43, block=5) + q, main_kv, block_kv, kv_cache, slots, start_pos, sink, _, freqs = _make_inputs(43, block=block) draft_block, slots_i32, cache_seqs = _prepare_attention_inputs( main_kv, block_kv, kv_cache, slots, start_pos ) @@ -630,7 +691,8 @@ def test_attention_cache_miss_rejects_cuda_graph_capture(monkeypatch): ) -def test_preparation_cache_misses_reject_cuda_graph_capture(monkeypatch): +@pytest.mark.parametrize("block", _PHYSICAL_KS) +def test_preparation_cache_misses_reject_cuda_graph_capture(monkeypatch, block): import tensorrt_llm._torch.custom_ops.dspark_rmsnorm_rope_custom_op as preparation_op preparation_op._compile_dspark_rmsnorm_rope_cache_write.cache_clear() @@ -645,10 +707,11 @@ def test_preparation_cache_misses_reject_cuda_graph_capture(monkeypatch): with pytest.raises(RuntimeError, match="cache-write must be warmed up"): preparation_op._compile_dspark_rmsnorm_rope_cache_write(1e-6) with pytest.raises(RuntimeError, match="draft-block must be warmed up"): - preparation_op._compile_dspark_rmsnorm_rope_draft_block(5, 1e-6) + preparation_op._compile_dspark_rmsnorm_rope_draft_block(block, 1e-6) -def test_preparation_self_jit_covers_dynamic_batches(): +@pytest.mark.parametrize("block", _PHYSICAL_KS) +def test_preparation_self_jit_covers_dynamic_batches(block): from tensorrt_llm._torch.custom_ops.dspark_rmsnorm_rope_custom_op import ( _compile_dspark_rmsnorm_rope_cache_write, _compile_dspark_rmsnorm_rope_draft_block, @@ -659,7 +722,7 @@ def test_preparation_self_jit_covers_dynamic_batches(): _compile_dspark_rmsnorm_rope_cache_write.cache_clear() _compile_dspark_rmsnorm_rope_draft_block.cache_clear() q, main_kv, block_kv, kv_cache, slots, start_pos, _, _, _ = _make_inputs( - 23, start_pos_values=[300, 4, 250], cache_pages=40 + 23, block=block, start_pos_values=[300, 4, 250], cache_pages=40 ) weight = torch.ones(512, device=q.device, dtype=q.dtype) main_freqs = torch.zeros(q.shape[0], _ROPE_DIM // 2, 2, device=q.device) @@ -668,7 +731,7 @@ def test_preparation_self_jit_covers_dynamic_batches(): block_freqs[..., 0] = 1 for batch in (1, 3, 32): - args = _make_inputs(31 + batch, batch=batch, cache_pages=40) + args = _make_inputs(31 + batch, block=block, batch=batch, cache_pages=40) q_b, main_b, block_b, cache_b, slots_b, pos_b, _, _, _ = args main_freqs_b = torch.zeros(batch, _ROPE_DIM // 2, 2, device=q.device) block_freqs_b = torch.zeros(batch * q_b.shape[1], _ROPE_DIM // 2, 2, device=q.device) diff --git a/tests/unittest/llmapi/test_dspark_workload_admission_config.py b/tests/unittest/llmapi/test_dspark_workload_admission_config.py new file mode 100644 index 000000000000..bf144525a4d0 --- /dev/null +++ b/tests/unittest/llmapi/test_dspark_workload_admission_config.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Startup-boundary tests for DSpark workload admission.""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from tensorrt_llm.llmapi.llm_args import CudaGraphConfig, DSparkDecodingConfig, TorchLlmArgs + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("outcome", ["missing", "negative", "stale"]) +def test_dspark_workload_admission_disables_before_runtime_construction(tmp_path, outcome): + (tmp_path / "config.json").write_text('{"dspark_block_size": 2}') + spec_cfg = DSparkDecodingConfig( + max_draft_len=2, + speculative_model=str(tmp_path), + enable_confidence_scheduling=True, + validate_confidence_artifact_freshness=True, + confidence_sps_table_path=(None if outcome == "missing" else "/sealed/sps.json"), + confidence_sps_live_fingerprint_path=( + None if outcome == "missing" else "/sealed/live.json" + ), + confidence_admission_receipt_path=( + None if outcome == "missing" else "/sealed/admission.json" + ), + confidence_admission_receipt_sha256=(None if outcome == "missing" else "a" * 64), + ) + if outcome == "negative": + patch_kwargs = {"return_value": SimpleNamespace(admitted=False)} + elif outcome == "stale": + patch_kwargs = {"side_effect": ValueError("stale receipt")} + else: + patch_kwargs = {"side_effect": TypeError("missing receipt")} + with patch( + "tensorrt_llm._torch.speculative.dspark_planner.load_confidence_workload_admission", + **patch_kwargs, + ): + # No CUDA-graph configuration is supplied. Construction succeeds only + # if admission resolves the feature off before the confidence-specific + # environment validator observes it. + args = TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + ) + + assert args.speculative_config.max_draft_len == 2 + assert args.speculative_config.enable_confidence_scheduling is False + assert args.speculative_config.confidence_sps_table_path is None + assert args.speculative_config.confidence_sps_live_fingerprint_path is None + assert args.speculative_config.confidence_admission_receipt_path is None + assert args.speculative_config.confidence_admission_receipt_sha256 is None + + +@pytest.mark.cpu_only +def test_dspark_positive_workload_admission_preserves_confidence_path(tmp_path): + (tmp_path / "config.json").write_text('{"dspark_block_size": 2}') + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"mtp.0.weight": "model.safetensors"}}' + ) + spec_cfg = DSparkDecodingConfig( + max_draft_len=2, + speculative_model=str(tmp_path), + enable_confidence_scheduling=True, + validate_confidence_artifact_freshness=True, + confidence_sps_table_path="/sealed/sps.json", + confidence_sps_live_fingerprint_path="/sealed/live.json", + confidence_admission_receipt_path="/sealed/admission.json", + confidence_admission_receipt_sha256="a" * 64, + ) + with patch( + "tensorrt_llm._torch.speculative.dspark_planner.load_confidence_workload_admission", + return_value=SimpleNamespace(admitted=True), + ): + args = TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + cuda_graph_config=CudaGraphConfig(batch_sizes=[1], enable_padding=True), + ) + + assert args.speculative_config.enable_confidence_scheduling is True + assert args.speculative_config.confidence_admission_receipt_path == "/sealed/admission.json" + assert args.speculative_config.confidence_admission_receipt_sha256 == "a" * 64 + + +@pytest.mark.cpu_only +def test_dspark_reuse_keeps_enabled_without_workload_admission(tmp_path): + (tmp_path / "config.json").write_text('{"dspark_block_size": 2}') + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"mtp.0.weight": "model.safetensors"}}' + ) + spec_cfg = DSparkDecodingConfig( + max_draft_len=2, + speculative_model=str(tmp_path), + enable_confidence_scheduling=True, + confidence_sps_table_path="/original/sps.json", + confidence_sts_path="/original/sts.json", + ) + with patch( + "tensorrt_llm._torch.speculative.dspark_planner.load_confidence_workload_admission", + side_effect=AssertionError("Reuse must not construct an admission receipt"), + ): + args = TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + cuda_graph_config=CudaGraphConfig(batch_sizes=[1], enable_padding=True), + ) + + assert args.speculative_config.enable_confidence_scheduling is True + assert args.speculative_config.validate_confidence_artifact_freshness is False + assert args.speculative_config.confidence_admission_receipt_path is None + assert args.speculative_config.confidence_admission_receipt_sha256 is None + assert args.speculative_config.confidence_sps_table_path == "/original/sps.json" + assert args.speculative_config.confidence_sts_path == "/original/sts.json" + + +@pytest.mark.cpu_only +def test_dspark_reuse_preserves_padded_graph_requirement(tmp_path): + (tmp_path / "config.json").write_text('{"dspark_block_size": 2}') + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"mtp.0.weight": "model.safetensors"}}' + ) + spec_cfg = DSparkDecodingConfig( + max_draft_len=2, + speculative_model=str(tmp_path), + enable_confidence_scheduling=True, + confidence_sps_table_path="/original/sps.json", + ) + with pytest.raises(ValueError, match="cuda_graph_config"): + TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + cuda_graph_config=None, + ) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index dd0e9298d4d4..0756a56c066d 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -384,7 +384,7 @@ def test_llm_args_with_kvcache_config(self): llm_args_dict = self._yaml_to_dict(yaml_content) llm_args = TorchLlmArgs(model=llama_model_path, **llm_args_dict) - assert llm_args.kv_cache_config.enable_block_reuse == True + assert llm_args.kv_cache_config.enable_block_reuse assert llm_args.kv_cache_config.max_tokens == 1024 assert llm_args.kv_cache_config.max_attention_window == [ 1024, 1024, 1024 @@ -1198,7 +1198,7 @@ def test_KvCacheConfig_declaration(): attention_dp_events_gather_period_ms=10) pybind_config = config._to_pybind() - assert pybind_config.enable_block_reuse == True + assert pybind_config.enable_block_reuse assert pybind_config.max_tokens == 1024 assert pybind_config.max_attention_window == [1024, 1024, 1024] assert pybind_config.free_gpu_memory_fraction == 0.5 @@ -1234,8 +1234,8 @@ def test_KvCacheConfig_declaration(): ).kv_cache_event_hash_algo == "v2_sha256" assert KvCacheConfig(kv_cache_event_hash_algo="v2_sha256_64" ).kv_cache_event_hash_algo == "v2_sha256_64" - assert pybind_config.enable_partial_reuse == True - assert pybind_config.copy_on_partial_reuse == True + assert pybind_config.enable_partial_reuse + assert pybind_config.copy_on_partial_reuse assert pybind_config.attention_dp_events_gather_period_ms == 10 assert BlockReuseConfig( policy="per_conversation").policy == "per_conversation" @@ -1864,8 +1864,8 @@ def test_DynamicBatchConfig_declaration(): pybind_config = PybindMirror.maybe_to_pybind(config) - assert pybind_config.enable_batch_size_tuning == True - assert pybind_config.enable_max_num_tokens_tuning == True + assert pybind_config.enable_batch_size_tuning + assert pybind_config.enable_max_num_tokens_tuning assert pybind_config.dynamic_batch_moving_average_window == 10 @@ -3112,7 +3112,7 @@ def test_dynamic_batch_config_arbitrary_args(self): config = DynamicBatchConfig(enable_batch_size_tuning=True, enable_max_num_tokens_tuning=True, dynamic_batch_moving_average_window=8) - assert config.enable_batch_size_tuning == True + assert config.enable_batch_size_tuning # Arbitrary arguments should be rejected with pytest.raises( @@ -3157,7 +3157,7 @@ def test_kv_cache_config_arbitrary_args(self): """Test that KvCacheConfig rejects arbitrary arguments.""" # Valid arguments should work config = KvCacheConfig(enable_block_reuse=True, max_tokens=1024) - assert config.enable_block_reuse == True + assert config.enable_block_reuse assert config.max_tokens == 1024 # Arbitrary arguments should be rejected @@ -3172,8 +3172,8 @@ def test_extended_runtime_perf_knob_config_arbitrary_args(self): # Valid arguments should work config = ExtendedRuntimePerfKnobConfig(multi_block_mode=True, cuda_graph_mode=False) - assert config.multi_block_mode == True - assert config.cuda_graph_mode == False + assert config.multi_block_mode + assert not config.cuda_graph_mode # Arbitrary arguments should be rejected with pytest.raises( @@ -3288,8 +3288,8 @@ def test_torch_compile_config_arbitrary_args(self): # Valid arguments should work config = TorchCompileConfig(enable_fullgraph=True, enable_inductor=False) - assert config.enable_fullgraph == True - assert config.enable_inductor == False + assert config.enable_fullgraph + assert not config.enable_inductor # Arbitrary arguments should be rejected with pytest.raises( @@ -3775,8 +3775,8 @@ def test_model_defaults_application(self): assert self.get_model_defaults_called modified_args = llm._executor.engine.model_engine.llm_args - assert modified_args.enable_chunked_prefill == True - assert modified_args.kv_cache_config.enable_block_reuse == False + assert modified_args.enable_chunked_prefill + assert not modified_args.kv_cache_config.enable_block_reuse assert modified_args.kv_cache_config.free_gpu_memory_fraction == 0.75 @pytest.mark.part0 @@ -3796,10 +3796,10 @@ def test_user_overrides_respected(self): assert self.get_model_defaults_called modified_args = llm._executor.engine.model_engine.llm_args - assert modified_args.enable_chunked_prefill == False + assert not modified_args.enable_chunked_prefill assert modified_args.max_batch_size == 42 assert modified_args.max_input_len == 256 - assert modified_args.kv_cache_config.enable_block_reuse == True + assert modified_args.kv_cache_config.enable_block_reuse @pytest.mark.part0 def test_partial_user_override(self): @@ -3816,8 +3816,8 @@ def test_partial_user_override(self): modified_args = llm._executor.engine.model_engine.llm_args assert modified_args.max_batch_size == 42 - assert modified_args.enable_chunked_prefill == True - assert modified_args.kv_cache_config.enable_block_reuse == False + assert modified_args.enable_chunked_prefill + assert not modified_args.kv_cache_config.enable_block_reuse @pytest.mark.part0 def test_empty_nested_config_preserves_defaults(self): @@ -3843,7 +3843,7 @@ def test_empty_nested_config_preserves_defaults(self): # Model defaults set enable_block_reuse=False and # free_gpu_memory_fraction=0.75. An empty KvCacheConfig() # should not prevent these from being applied. - assert modified_args.kv_cache_config.enable_block_reuse == False + assert not modified_args.kv_cache_config.enable_block_reuse assert modified_args.kv_cache_config.free_gpu_memory_fraction == 0.75 @@ -3914,8 +3914,7 @@ def _get_qualified_name(cls: type) -> str: @pytest.mark.cpu_only class TestPydanticBestPractices: - """Ensure that the user-facing LlmArgs and its subfields follow Pydantic best practices. - """ + """Check user-facing LlmArgs and subfields for Pydantic best practices.""" # Fields exempt from Pydantic compatibility checks due to typing limitations or other edge cases. # Avoid adding to this list unless absolutely necessary, especially if a field is user-facing. @@ -5260,3 +5259,154 @@ def test_deepseek_resolves_auto_to_python_on_nixl(self) -> None: cfg = self._pretrained_config(["DeepseekV3ForCausalLM"], "deepseek_v3") _resolve_transceiver_runtime_auto(args, DeepseekV3ForCausalLM, cfg) assert args.cache_transceiver_config.transceiver_runtime == "PYTHON" + + +class TestDSparkConfidenceScheduling: + """Config surface for DSpark confidence-scheduled verification.""" + + def _cfg(self, **kw): + from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig + kw.setdefault("max_draft_len", 5) + return DSparkDecodingConfig(**kw) + + def _scheduled(self, **kw): + """Scheduling on with the required profiled cost table.""" + kw["enable_confidence_scheduling"] = True + kw.setdefault("confidence_sps_table_path", "/tmp/sps.json") + kw.setdefault("confidence_sps_live_fingerprint_path", + "/tmp/runtime.json") + return self._cfg(**kw) + + def test_defaults_are_off(self): + c = self._cfg() + assert c.enable_confidence_scheduling is False + assert c.enable_fused_confidence_scheduler is False + + def test_knobs_require_the_master_switch(self): + for kw in ({ + "confidence_sts_path": "/tmp/sts.json" + }, { + "confidence_sps_table_path": "/tmp/sps.json" + }, { + "confidence_sps_live_fingerprint_path": "/tmp/runtime.json" + }, { + "confidence_admission_receipt_path": "/tmp/admission.json" + }, { + "confidence_admission_receipt_sha256": "a" * 64 + }): + with pytest.raises(ValueError, + match="enable_confidence_scheduling"): + self._cfg(**kw) + + def test_admission_artifacts_are_resolved_at_torch_args_boundary(self): + c = self._cfg(enable_confidence_scheduling=True) + assert c.enable_confidence_scheduling is True + + def test_removed_legacy_verify_len_tiers_are_rejected(self): + with pytest.raises(ValidationError, + match="confidence_verify_len_tiers"): + self._scheduled(confidence_verify_len_tiers=[1, 3, 5]) + + def test_physical_k1_resolves_confidence_scheduling_to_static(self): + c = self._scheduled( + max_draft_len=1, + confidence_sts_path="/tmp/sts.json", + confidence_admission_receipt_path="/tmp/admission.json", + confidence_admission_receipt_sha256="a" * 64, + enable_fused_confidence_scheduler=True, + ) + + assert c.enable_confidence_scheduling is False + assert c.enable_fused_confidence_scheduler is False + assert c.confidence_sts_path is None + assert c.confidence_sps_table_path is None + assert c.confidence_sps_live_fingerprint_path is None + assert c.confidence_admission_receipt_path is None + assert c.confidence_admission_receipt_sha256 is None + + def test_physical_k1_needs_no_confidence_artifacts(self): + c = self._cfg(max_draft_len=1, enable_confidence_scheduling=True) + + assert c.enable_confidence_scheduling is False + assert c.max_total_draft_tokens == 1 + assert c.tokens_per_gen_step == 2 + + def test_physical_k1_remains_valid_when_scheduling_is_off(self): + c = self._cfg(max_draft_len=1) + assert c.enable_confidence_scheduling is False + assert c.max_draft_len == 1 + assert c.max_total_draft_tokens == 1 + assert c.tokens_per_gen_step == 2 + + def test_fused_scheduler_is_independently_gated(self): + with pytest.raises(ValueError, + match="enable_fused_confidence_scheduler"): + self._cfg(enable_fused_confidence_scheduler=True) + c = self._scheduled(enable_fused_confidence_scheduler=True) + assert c.enable_fused_confidence_scheduler is True + + @staticmethod + def _environment(**overrides): + values = dict( + cuda_graph_config=SimpleNamespace(enable_padding=True), + speculative_config=SimpleNamespace( + draft_is_embedded_in_target=True), + guided_decoding_backend=None, + enable_lora=False, + lora_config=None, + sparse_attention_config=None, + pipeline_parallel_size=1, + ) + values.update(overrides) + return SimpleNamespace(**values) + + def test_confidence_environment_rejects_standalone_dspark(self): + environment = self._environment(speculative_config=SimpleNamespace( + draft_is_embedded_in_target=False)) + with pytest.raises(ValueError, + match="only the embedded DeepSeek-V4 draft"): + TorchLlmArgs._validate_dspark_confidence_environment(environment) + + @pytest.mark.parametrize( + ("overrides", "message"), + [ + ({ + "cuda_graph_config": None + }, "requires cuda_graph_config"), + ( + { + "cuda_graph_config": SimpleNamespace(enable_padding=False) + }, + "enable_padding=True", + ), + ({ + "guided_decoding_backend": "xgrammar" + }, "guided decoding"), + ({ + "enable_lora": True + }, "does not yet support LoRA"), + ({ + "pipeline_parallel_size": 2 + }, "does not yet support pipeline parallelism"), + ( + { + "sparse_attention_config": + SimpleNamespace( + enable_heuristic_topk=True, + use_cute_dsl_topk=False, + use_cute_dsl_paged_mqa_logits=False, + ) + }, + "enable_heuristic_topk", + ), + ], + ) + def test_confidence_environment_rejects_unsupported_runtime_combinations( + self, overrides, message): + environment = self._environment(**overrides) + with pytest.raises(ValueError, match=message): + TorchLlmArgs._validate_dspark_confidence_environment(environment) + + def test_confidence_environment_accepts_padded_cuda_graphs(self): + TorchLlmArgs._validate_dspark_confidence_environment( + self._environment()) From ff827dbda4ea875da741abda14190341210f735f Mon Sep 17 00:00:00 2001 From: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> Date: Thu, 15 Oct 2026 11:58:16 -0700 Subject: [PATCH 3/7] [TRTLLM-14620][fix] Preserve measured-table reuse for configured graph subsets Retain the original complete payload and fingerprint validation while restricting installed runtime choices to captured graph tiers. Strict freshness keeps exact ladder matching; drain guard requires both captured and measured tail graphs. Local source-only publication checkpoint; native qualification remains separate. Signed-off-by: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> --- .../_torch/speculative/dspark_planner.py | 33 +++++++++++++------ 1 file changed, 23 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dspark_planner.py b/tensorrt_llm/_torch/speculative/dspark_planner.py index 0d87e590ab4b..e9c6222ba8fe 100644 --- a/tensorrt_llm/_torch/speculative/dspark_planner.py +++ b/tensorrt_llm/_torch/speculative/dspark_planner.py @@ -382,7 +382,9 @@ def load_runtime_sps_cost_table( """Load structurally checked exact costs, optionally checking freshness. Original artifact metadata is returned unchanged. Reuse does not assert - that recorded timings still describe the active source or GPU node. + that recorded timings still describe the active source or GPU node. Only + configured graph tiers are installed as runtime execution choices; unused + measured tiers remain in the original returned payload. """ payload = _read_sps_cost_payload(path) schema_version = payload.get("schema_version") @@ -406,7 +408,8 @@ def load_runtime_sps_cost_table( validate_freshness=validate_freshness, ) return _build_exact_sps_cost_table( - payload, max_draft_len=max_draft_len, validate_artifact_freshness=validate_freshness + payload, max_draft_len=max_draft_len, validate_artifact_freshness=validate_freshness, + runtime_graph_batch_sizes=graph_batch_sizes, ), payload @@ -415,12 +418,18 @@ def _build_exact_sps_cost_table( *, max_draft_len: int, validate_artifact_freshness: bool = True, + runtime_graph_batch_sizes: Sequence[int] | None = None, ) -> ExactSpsCostTable: exact_tables = payload["cost_tables"] assert isinstance(exact_tables, dict) parsed_tables: dict[int, ExactSpsCostRow] = {} + # Scope execution choices, not serialized measurements. The loader already + # validates the entire original artifact, including unused finite cells. + runtime_graphs = set(runtime_graph_batch_sizes) if runtime_graph_batch_sizes is not None else None for graph_batch_size, exact_payload in exact_tables.items(): canonical_graph_batch_size = _parse_graph_batch_size_key(graph_batch_size) + if runtime_graphs is not None and canonical_graph_batch_size not in runtime_graphs: + continue assert isinstance(exact_payload, dict) parsed_tables[canonical_graph_batch_size] = ExactSpsCostRow( token_counts=tuple( @@ -737,7 +746,8 @@ def validate_sps_cost_table_payload( Raises: TypeError: If the artifact has malformed container types. ValueError: If its schema, fingerprint, coverage, or provenance does - not exactly describe the configured graph ladder. + not cover the configured graph ladder. Strict freshness additionally + requires the original full graph ladder and exact live identity. """ if not isinstance(payload, dict): raise TypeError("SPS cost artifact must contain a JSON object") @@ -802,19 +812,22 @@ def validate_sps_cost_table_payload( field="iteration drain tail_graph_batch_size", minimum=1, ) - if tail_graph_batch_size not in measured_tables: + if (tail_graph_batch_size not in measured_tables + or tail_graph_batch_size not in configured_graph_batch_sizes): raise ValueError( - "Iteration drain guard requires a measured native tail table for " + "Iteration drain guard requires a captured, measured native tail table for " f"G={tail_graph_batch_size}" ) fingerprint_graph_batch_sizes = set(fingerprint["rank_local_graph_batch_sizes"]) if ( - configured_graph_batch_sizes != set(measured_tables) - or configured_graph_batch_sizes != fingerprint_graph_batch_sizes + set(measured_tables) != fingerprint_graph_batch_sizes + or not configured_graph_batch_sizes.issubset(measured_tables) + or (validate_freshness and configured_graph_batch_sizes != set(measured_tables)) ): raise ValueError( - "Multi-G SPS graph batch sizes must match exactly across runtime " - "graphs, cost_tables, and engine_fingerprint" + "Runtime graphs require measured coverage; cost_tables and their " + "original engine_fingerprint must match exactly. Strict freshness " + "also requires the original runtime graph ladder." ) runtime_max_draft_len = _require_json_int( max_draft_len, field="runtime max_draft_len", minimum=1 @@ -900,7 +913,7 @@ def validate_sps_cost_table_payload( gpu_count = fingerprint["gpu_count"] global_graph_batch_sizes = set(fingerprint["global_graph_batch_sizes"]) if global_graph_batch_sizes != { - graph_batch_size * gpu_count for graph_batch_size in configured_graph_batch_sizes + graph_batch_size * gpu_count for graph_batch_size in measured_tables }: raise ValueError( "SPS engine_fingerprint global_graph_batch_sizes are inconsistent " From 48799289924f1db8d7a0884458ca74ed833621ed Mon Sep 17 00:00:00 2001 From: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> Date: Thu, 15 Oct 2026 12:12:07 -0700 Subject: [PATCH 4/7] [TRTLLM-14620][feat] Add reusable DSpark measurement and reproduction tooling Local unpublished tooling candidate stacked on the reviewed PR5 and reconciled PR6 plus configured-graph-subset reuse. Preserve ordinary default construction and explicit offline profiling/observation; forward replay position offsets in both adapters with permanent regression coverage. CPU/source checks only; matching-runtime pytest, native/GPU qualification and schema readiness remain separate and unperformed. Full Windows hook invocation remains incomplete due recorded vendor line-ending and script-interpreter limitations. Signed-off-by: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> --- .../deepseek_v4/CONFIDENCE_REPRODUCTION.md | 187 ++++++ .../models/core/deepseek_v4/COST_EXPORT.md | 128 ++++ .../models/core/deepseek_v4/COST_PROFILING.md | 289 +++++++++ examples/models/core/deepseek_v4/README.md | 17 +- tensorrt_llm/_torch/pyexecutor/_util.py | 4 +- .../pyexecutor/dspark_observation_runtime.py | 341 +++++++++++ .../pyexecutor/dspark_profile_engine.py | 462 ++++++++++++++ .../dspark_profile_graph_storage.py | 188 ++++++ .../pyexecutor/dspark_request_profile.py | 535 +++++++++++++++++ .../engine/runners/decoder/runner.py | 10 +- .../_torch/pyexecutor/py_executor_creator.py | 32 +- .../_torch/speculative/dspark_cost_profile.py | 90 +++ .../speculative/dspark_cost_profiler.py | 563 ++++++++++++++++++ .../_torch/speculative/dspark_reproduction.py | 370 ++++++++++++ tensorrt_llm/executor/base_worker.py | 11 +- tensorrt_llm/tools/dspark_calibration.py | 239 ++++++++ tensorrt_llm/tools/dspark_cost_export.py | 413 +++++++++++++ tensorrt_llm/tools/dspark_observations.py | 213 +++++++ tensorrt_llm/tools/dspark_reproduction.py | 73 +++ tensorrt_llm/tools/verify_costs.py | 445 ++++++++++++++ .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../test_dspark_reproduction_factory.py | 397 ++++++++++++ .../unittest/tools/test_dspark_calibration.py | 156 +++++ .../unittest/tools/test_dspark_cost_export.py | 409 +++++++++++++ .../tools/test_dspark_observations.py | 445 ++++++++++++++ 25 files changed, 6004 insertions(+), 14 deletions(-) create mode 100644 examples/models/core/deepseek_v4/CONFIDENCE_REPRODUCTION.md create mode 100644 examples/models/core/deepseek_v4/COST_EXPORT.md create mode 100644 examples/models/core/deepseek_v4/COST_PROFILING.md create mode 100644 tensorrt_llm/_torch/pyexecutor/dspark_observation_runtime.py create mode 100644 tensorrt_llm/_torch/pyexecutor/dspark_profile_engine.py create mode 100644 tensorrt_llm/_torch/pyexecutor/dspark_profile_graph_storage.py create mode 100644 tensorrt_llm/_torch/pyexecutor/dspark_request_profile.py create mode 100644 tensorrt_llm/_torch/speculative/dspark_cost_profile.py create mode 100644 tensorrt_llm/_torch/speculative/dspark_cost_profiler.py create mode 100644 tensorrt_llm/_torch/speculative/dspark_reproduction.py create mode 100644 tensorrt_llm/tools/dspark_calibration.py create mode 100644 tensorrt_llm/tools/dspark_cost_export.py create mode 100644 tensorrt_llm/tools/dspark_observations.py create mode 100644 tensorrt_llm/tools/dspark_reproduction.py create mode 100644 tensorrt_llm/tools/verify_costs.py create mode 100644 tests/unittest/executor/test_dspark_reproduction_factory.py create mode 100644 tests/unittest/tools/test_dspark_calibration.py create mode 100644 tests/unittest/tools/test_dspark_cost_export.py create mode 100644 tests/unittest/tools/test_dspark_observations.py diff --git a/examples/models/core/deepseek_v4/CONFIDENCE_REPRODUCTION.md b/examples/models/core/deepseek_v4/CONFIDENCE_REPRODUCTION.md new file mode 100644 index 000000000000..37e277451747 --- /dev/null +++ b/examples/models/core/deepseek_v4/CONFIDENCE_REPRODUCTION.md @@ -0,0 +1,187 @@ +# Confidence calibration and a fair static/enabled comparison + +This is a **source/CPU-reviewed candidate workflow**. These commands have not yet +completed native full-model GPU validation. Do not describe this document as an +already-supported public reproduction recipe until that validation and +publication finish. Its tools have no private harness or cluster-name/path, +Slurm account, fixed K3 or fixed BS512 binding. + +The two routes are distinct: + +- Reuse a structurally valid supplied calibration/cost table, preserving their + original measurements and provenance. Once confidence is explicitly enabled, + reuse is the default and warnings explain that predictions may be stale. +- Generate new calibration and measured whole-step costs using the offline + commands below. Optional freshness/economic validation is a separate choice; + it is not required before reusing valid artifacts. + +Confidence scheduling remains off by default. Target verification, sampling, +KV/graph contracts, kernel guards and fatal errors are not relaxed by reuse. +Poor cost/calibration estimates can choose a slower verifier shape; reuse is +not a performance guarantee. + +## 1. Build or reuse an ordinary compatible runtime + +Start from a complete public checkout containing these tools, the configurable +cost-profile adapter, and the default-reuse consumer. Use the repository's +[source-build instructions](../../../../docs/source/installation/build-from-source.md) +or its supported precompiled/editable installation route for compatible native +artifacts. Unrelated Python changes do not force native recompilation. Preserve +the original native provenance; the tools separately record loaded Python bytes. + +Use your cluster's ordinary allocation/container/MPI launch procedure and model +mounts. Keep the intended precision, attention backend, fused/prenorm paths, +parallelism and original startup/model/shutdown budget. Do not change numerical +configuration or disable failed hardware/operator checks to make a command run. +The code neither repairs machines nor changes cluster/security policy. + +The complete constructor controls run without constructing a model, inside the +actual driver-backed installed runtime before a costly model caller. Some +native imports require an available driver even though these controls do not +execute model or GPU work; do not substitute a driverless import failure: + +```bash +python -m pytest -q tests/unittest/executor/test_dspark_reproduction_factory.py +``` + +This exercises the actual public argument validator and the existing internal +factory seam without constructing a model. Internal `executor_cls` is not a +public LLM argument and remains rejected. Passing this check is not a native +GPU/operator/model result. + +Prepare tokenized JSONL inputs: each line has a nonempty integer +`prompt_token_ids` list. Model weights, inputs and raw observations stay in +access-controlled storage; they are not published with the tooling. The selected +first `--batch-size` rows form one fixed request cohort. Each request generates +the same requested output length with greedy sampling and EOS ignored. The +original sequence-length limit must contain prompt plus output. + +Use a YAML accepted by `TorchLlmArgs`: matching physical K, intended topology, +enough rank-local request/token/KV capacity and padded decode graphs. Aggregate +PP=1/CP=1 generation is supported by this initial driver. Guided decoding, +LoRA, custom logits processing, disaggregated transfer, in-graph sampling and +resource-governor control are not part of this bounded fixed-cohort workflow. +No alternative attention backend is selected by the tool. + +## 2. Reuse supplied artifacts, if that is your choice + +Set `--cost-table` and optionally `--calibration` for the dynamic command in +step 5. No new cost profiling, forged live fingerprint or admission receipt is +needed first. The normal loader still validates table structure, physical K, +legal captured graph cells, finite costs and calibration geometry. Missing or +incompatible cost coverage cannot be reported as dynamic throughput. + +`validate_confidence_artifact_freshness` stays at its public default `False`. +Users who choose `True` must supply the corresponding genuine destination +fingerprint and workload-admission artifacts; this driver does not fabricate +them. Original cost/calibration hashes are recorded separately from the runtime +actually used for a new comparison. + +## 3. Generate calibration from real native full-K proposals + +```bash +mpirun -n "$RANKS" python tensorrt_llm/tools/dspark_reproduction.py run \ + --mode collect --model "$MODEL_DIR" --config "$RUNTIME_YAML" \ + --inputs "$INPUTS_JSONL" --source-root "$SOURCE_ROOT" \ + --physical-k "$K" --batch-size "$BS" --output-tokens "$OUTPUT_TOKENS" \ + --timeout "$FULL_BUDGET_SECONDS" --output "$RUN_ROOT/calibration-observations" + +python tensorrt_llm/tools/dspark_calibration.py \ + --shards "$RUN_ROOT"/calibration-observations/rank-*-pairs.pt \ + --physical-k "$K" --output "$RUN_ROOT/calibration.json" +``` + +The current collector supports the original overlap scheduler input seam; +`disable_overlap_scheduler: False` must already be your intended configuration. +It runs native full-K verification and evaluates the existing confidence head. +There is no confidence-selected shortening and no loaded cost table. Real +warm-up cohorts may contribute calibration observations, but never benchmark +timing. The fitter's default sample floor is 1000 genuine paired rows; collect +enough requests/output tokens rather than weakening that floor for real use. + +Pairing is not by arrival order. Step n's confidence and newly proposed tokens +describe the proposal verified later, usually at n+1. The adapter binds that +later forward to the exact original input `SampleState.device` object. It +snapshots verification-input proposals before persistent buffers can be reused, +and captures the producing confidence row and draft stamp before the normal +sampler records its existing event. It reads those host copies only after the +original sampler's normal synchronization/update. Request incarnation, source +proposal content, worker-local confidence row, sampler slot and full-K extent +must agree. Prefill placeholders, retired requests and unscored bootstrap rows +are counted, not fitted as genuine observations. Wrong-step/stale/censored +pairs fail; the current live confidence is never substituted. + +Expected files are `rank-N-pairs.pt`, `rank-N-run.json` and the fitted JSON with +`sts_temperatures`/`temperatures`. Each shard has logits `[rows,K]`, cumulative +accepted-prefix labels, exact join metadata and actual source/runtime identity. +These observations do not establish downstream task accuracy by themselves. + +## 4. Generate directly measured costs + +Follow [Offline cost profiling](COST_PROFILING.md) and +[Whole-step export](COST_EXPORT.md). Declare native V0 plus legal compact cells +for your configurable physical K and resolved rank-local graph tiers. Use the +`--cost-scope whole-step` producer and export its actual completed measurements +to the schema-v2 table consumed by `confidence_sps_table_path`. + +Do not substitute graph-only timing, sum overlapping components or invent +serving prices from calibration samples. The whole-step producer reuses the +original executor-loop timing boundary; profile-only forced cells cannot serve +ordinary traffic or prove automatic-policy gains. Keep measured confidence-head +and policy cost boundaries explicit. This candidate benchmark measures actual +online policy in its E2E time; it does not manufacture a separate overhead or +economic-admission receipt from that result. + +For attention DP on, global BS determines rank-local G through the attention-DP +fanout. The current schema-v2 cost exporter supports this fanout or one rank; +multi-rank attention-DP-off costs need their separate schema decision. This +restriction belongs to this cost export, not to generic DSpark support. + +## 5. Benchmark ordinary static and genuinely enabled execution + +Each command is an independent model lifetime. Use the same original inputs, +physical K, runtime YAML, output work, warm-up count and full feasible budget: + +```bash +mpirun -n "$RANKS" python tensorrt_llm/tools/dspark_reproduction.py run \ + --mode static --model "$MODEL_DIR" --config "$RUNTIME_YAML" \ + --inputs "$INPUTS_JSONL" --source-root "$SOURCE_ROOT" \ + --physical-k "$K" --batch-size "$BS" --output-tokens "$OUTPUT_TOKENS" \ + --timeout "$FULL_BUDGET_SECONDS" --output "$RUN_ROOT/static" + +mpirun -n "$RANKS" python tensorrt_llm/tools/dspark_reproduction.py run \ + --mode dynamic --model "$MODEL_DIR" --config "$RUNTIME_YAML" \ + --inputs "$INPUTS_JSONL" --source-root "$SOURCE_ROOT" \ + --physical-k "$K" --batch-size "$BS" --output-tokens "$OUTPUT_TOKENS" \ + --cost-table "$COST_TABLE_JSON" --calibration "$CALIBRATION_JSON" \ + --timeout "$FULL_BUDGET_SECONDS" --output "$RUN_ROOT/dynamic" + +python tensorrt_llm/tools/dspark_reproduction.py compare \ + --runs "$RUN_ROOT/static/rank-0-run.json" "$RUN_ROOT/dynamic/rank-0-run.json" \ + --output "$RUN_ROOT/comparison.json" +``` + +This uses the normal worker/request/executor/automatic confidence policy, not +the prescribed profile-only engine. Benchmark arms make no calibration copies. +Actual completed decode graph keys are recorded. Warm-up routing is reset before +the retained cohort; requested flags alone do not prove enabled execution. +The comparison requires both resolved config and runtime enabled status, actual +compact replay in the retained dynamic cohort, unchanged numerical runtime/input +work, completed outputs and natural shutdown. An enabled policy that only chose +native is a valid execution, but this consumer refuses to label it a reproduced +compact/dynamic gain. + +Reported TPS is global fixed-cohort enqueue-to-final-response throughput, +including prefill and final drain. It is not measured TPOT, sustained-refill +throughput, isolated verifier time or the historical plateau measurement. +No confidence interval is fabricated. Repeat in balanced order with fresh +lifetimes if you need a stable estimate; feed equal numbers of static and dynamic +records to `compare`. Exact output-token equality is reported separately and is +not silently substituted for the model's original task-accuracy tests. + +All previously tried batch sizes 128/256/512/1024/1536/2048 and source-backed +smaller points use the same CLI, with compatible configured graph/K/cost coverage. +K is an input, not fixed to K3. K1 native calibration/static execution is legal +but cannot demonstrate compact dynamic scheduling. There is no new broad BS +validation sweep implied: existing and freshly tested points must be labeled +truthfully, with unperformed points listed as prepared, not validated. diff --git a/examples/models/core/deepseek_v4/COST_EXPORT.md b/examples/models/core/deepseek_v4/COST_EXPORT.md new file mode 100644 index 000000000000..0368d0e645c6 --- /dev/null +++ b/examples/models/core/deepseek_v4/COST_EXPORT.md @@ -0,0 +1,128 @@ + + + +# Handing measured costs to DSpark + +The offline graph collector and runtime cost consumer have different contracts. +The collector's `verify-cost-measured-summary-v1` output measures the original +graph replay entry, including static input copies. The consumer's schema-v2 +`step_time_ms` represents direct whole-step costs. Do not rename the replay +median to `step_time_ms`, sum potentially overlapping graph/host components, or +construct a new engine fingerprint from the partial loaded-runtime observation. + +The explicit `whole-step` profile mode reuses the executor's existing +`hostStepTimeMS` boundary. It joins the just-finished complete loop to its actual +executor iteration; it does not sum graph and host spans, introduce new CUDA +events, alter streams or add synchronization. The raw exporter retains all-rank +joins, warm-up exclusions and shutdown closure. Graph timing is retained only as +a separately labeled diagnostic in this mode. + +`dspark_cost_export.py` maps those direct whole-step medians into the existing +runtime schema without filling missing cells. It also accepts an existing +schema-v2 artifact and can copy its original bytes after structural validation. +Both paths use the actual runtime schema validator. Plain graph summaries remain +ineligible; that does not prevent reuse of structurally valid supplied costs. + +These are source/CPU-prepared development interfaces. The complete new GPU +collection/export workflow has not been exercised or qualified yet. + +For a previously measured cost table and a runtime capturing G16 and G64: + +```bash +python tensorrt_llm/tools/dspark_cost_export.py \ + --input /path/to/recorded-costs.json \ + --physical-k 3 --graph-batch-sizes 16 64 +``` + +`K` is configurable and has no fixed pilot restriction. `G` is the rank-local +request count of the actual captured graph, not always the global batch size. +With eight attention-DP ranks, BS128 and BS512 correspond to G16 and G64; with +attention DP off they do not. Pass the exact graph ladder the engine resolves. +Native V0, positive finite timing values, physical K, cell provenance and graph +coverage must satisfy the existing consumer. Missing cells are not interpolated, +copied from another K or filled with zero. No cluster name or site path is used. + +Add `--output /path/to/new-costs.json` only to deliver a new byte-identical copy. +The destination must not exist. No original values, recorded source head or +engine fingerprint are rewritten. In this reuse path there is no exact-head +freshness check and no economic admission requirement; the result warns that +destination freshness and performance are unchecked. Confidence scheduling is +still default-off and must be explicitly enabled through its normal configuration. + +The structural-check report is not a calibration result, GPU health check, +task-accuracy certificate, positive admission receipt or benchmark. Genuine +enabled routing, output checks and a matched static/dynamic comparison remain +separate measurements. The source-only CPU tests do not qualify this complete +clean-checkout procedure on a GPU. + +## Preparing new direct whole-step costs + +Use the configurable K/G/V grid and authorized inputs described in +[COST_PROFILING.md](COST_PROFILING.md). Set `enable_iter_perf_stats: true` in the +ordinary runtime YAML so the existing complete-loop timer is active. Use a +read-only checkout of the original source as `--source-root`; its actual Git +head/diff and loaded module bytes are recorded as provenance, not used as a +destination freshness gate. Runtime/library/GPU observations are preserved in +the raw identity. No new native build is implied by this Python-only tooling. +The Git head is only the original base: the raw `loaded_source_sha256` roster +and recorded runtime snapshot retain executed local Python overlays, including +untracked tooling. A dirty execution is not described as pristine just because +it has the same base head; original tracked diff bytes are hashed separately. + +The additional arguments to the normal profile interface are: + +```bash +python -m tensorrt_llm.tools.verify_costs profile \ + --model /path/to/authorized/checkpoint --config /path/to/runtime.yaml \ + --grid /path/to/profile-grid.json --physical-k 3 \ + --cost-scope whole-step --source-root /path/to/original-public-checkout \ + --inputs /path/to/tokenized-prompts.jsonl \ + --output /path/to/new-private-profile-directory +``` + +Use this command under the site's normal MPI/container launcher with the actual +GPU count and full startup/generation/export/shutdown budget. No cluster or +account is embedded in the tooling. This source-prepared interface is **not** a +GPU-tested reproduction instruction. No extra sweep is needed before reusing +valid supplied cost/calibration predictions. + +Reduce the closed raw observations with the existing `verify_costs.py export` +command. Its defaults retain at least 12 exact joins in each retained generation +and 24 pooled per cell. Then use the DSpark adapter: + +```bash +python tensorrt_llm/tools/dspark_cost_export.py \ + --input /path/to/whole-step-measured-summary.json \ + --physical-k 3 --graph-batch-sizes 16 64 \ + --output /path/to/new-costs.json +``` + +The adapter exports the retained per-iteration rank maximum's median unchanged; +every schema-v2 measurement references the original summary hash. The recorded +fingerprint is preserved, not refreshed by the CPU exporter. It neither creates +drain constants nor measures confidence-policy overhead, calibration quality or +economic admission. The default table policy margin remains 1%; it is not a +timing measurement. The original summary and raw observations remain available +to inspect scope, source/configuration and raw joins. + +The existing schema-v2 consumer checks global graph counts as rank-local G +times `gpu_count`. This new export currently supports attention-DP fanout or a +single rank; multi-rank attention-DP-off needs a separate schema decision and +is explicitly refused rather than recording incorrect global batch sizes. + +## Graph-only exports stay separate + +The default `graph-replay` collector records +`original-graph-replay-entry-including-static-copies` and +`DIAGNOSTIC_PARTIAL_LOADED_RUNTIME`, with +`production_sps_table_eligible: false`. This command deliberately refuses that +summary even if an eligibility flag is manually changed. Only choosing and +actually collecting the explicit whole-step boundary produces whole-step raw +observations. Changing a label does not extend an existing measured scope. + +GPU-forward values and overlap `iterLatencyMS` are not replacements for clean +complete-loop measurements: the former omits host work, while the latter can +span more than one loop. Optional policy-overhead/economic validation stays +separate. Authentic clean-checkout GPU collection, enabled-routing/output checks +and matched benchmarking are still required before publishing this as a verified +end-to-end self-generation recipe. diff --git a/examples/models/core/deepseek_v4/COST_PROFILING.md b/examples/models/core/deepseek_v4/COST_PROFILING.md new file mode 100644 index 000000000000..97a224862e2a --- /dev/null +++ b/examples/models/core/deepseek_v4/COST_PROFILING.md @@ -0,0 +1,289 @@ + + + +# Offline DSpark graph-cost profiling + +This is a development interface for collecting and exporting measured graph +replay costs. It is not yet a GPU-qualified colleague reproduction recipe. +The default graph-replay collector does not itself generate a serving SPS table, +paired calibration observations, an admission receipt, or an enabled throughput +result. The explicit complete-loop mode and CPU schema-v2 adapter are described +in [COST_EXPORT.md](COST_EXPORT.md); that new producer path remains GPU-unqualified. +The separate CPU calibration fitter below can fit already collected, correctly +paired shards. The complete public clean-checkout workflow is not yet qualified. + +The tools contain no private model, dataset, scratch path, historical receipt or +source-head pin. Use your own authorized local checkpoint and tokenized inputs. +Raw observations include prompt and output tokens; store them in an appropriately +access-controlled output directory. +The Linux collector creates only a new directory with mode0700 and new raw files +with mode0600, using exclusive/non-overwriting creation. It does not change +existing paths, ACLs, process umask or host settings. Existing/symlink targets are +rejected. POSIX permission behavior still needs actual Linux validation; Windows +CPU controls do not qualify Linux permissions or Windows ACLs. + +## Scope and terminology + +- `K` is the physical draft block produced by the model. Set it explicitly with + `--physical-k`; it must equal the runtime's `speculative_config.max_draft_len`. + Profiling never relabels a K5 execution as K3. +- `G` is the rank-local graph request count. +- `V` is the total verifier token budget in that rank-local graph. `V=0` denotes + the native full-block control, not an empty execution. Positive compact cells + require `K>1` and `G <= V < G*(K+1)`; K1 currently supports only native cells. +- The measured scope is **the original graph replay entry, including static + input copies**. It includes work already inside that graph, including the + embedded drafter and optional confidence head. It is not a target-only timer, + a complete confidence-policy price, or an end-to-end generation timer. + +The offline engine uses the original model loader, attention/backend selection, +graph capture/replay, request scheduler, KV allocator/estimator, sampler, token +updates and shutdown. It prescribes a declared verification cell; it does not +enable or bypass production confidence admission. Ordinary requests outside a +prescribed offline generation are rejected, and no serving RPC endpoint is exposed. +The profiler duplicates its control communicator collectively before worker +construction, retains it through shutdown/evidence closure, then frees it. +Original request broadcasts and runtime communicator/thread-local bindings are +unchanged. Multi-rank profiling requires existing MPI `THREAD_MULTIPLE` support; +the collector checks it and does not modify MPI initialization or thread level. + +Each cell runs a fixed cohort of real requests through prefill, generation and +drain, then repeats that generation. It does not time a single synthetic request +replayed in isolation. Only decode steps with exactly `G` requests on every rank +and the declared captured graph key become samples. Initial prefill, first-draft, +peer-not-ready and shorter drain batches retain the native path and are not timed +as that cell. No forced acceptance is allowed. The default reducer excludes +generation 0 and the first two eligible samples of every retained generation. + +## Prepare a matching runtime and workload + +Use a complete Linux checkout containing this tooling. Build/install TensorRT-LLM +using the repository's [build instructions](../../../../docs/source/installation/build-from-source.md), +or reuse genuinely compatible existing native artifacts through the documented +precompiled/editable-install route. A Python-only tooling change does not by +itself require a native rebuild or new cost/calibration measurements. +Do not combine these Python modules with incompatible native interfaces. Retain +the original native provenance and record changed Python source separately. +The collector records +the actual loaded Python/native-library hashes, package version, GPU class, +topology, graph configuration and observed attention-module classes. A module +roster is not proof of every dispatched kernel or of GPU health. +These raw identity observations are also not a complete weight/JIT/source/build +qualification or a production serving live fingerprint. + +Prepare a runtime YAML accepted by `TorchLlmArgs`, including: + +- `tensor_parallel_size` matching the MPI world size, `pipeline_parallel_size: 1` + and `context_parallel_size: 1`; +- the intended expert parallelism and `enable_attention_dp` setting; +- enough `max_batch_size`, `max_num_tokens`, `max_seq_len` and KV capacity for the + full cohort and requested output length; +- padded decode CUDA graphs whose **resolved rank-local ladder** exactly matches + all `G` values in the profile grid; +- an embedded DSpark configuration with the selected physical K, confidence + scheduling and the optional fused scheduler disabled, no draft-length schedule, + and no confidence SPS/fingerprint/admission/calibration fields. + +This initial driver is aggregate, fixed greedy generation. It rejects LoRA, +guided decoding, custom logits processors, disaggregated KV transfer/connectors, +encode-only execution and in-graph sampling. Prefill graph capture is not part of +this interface; resource-governor control is also unsupported. Keep the intended +numerical/quantization/backend configuration; +do not change it merely to make profiling launch. + +For example, an attention-DP eight-rank profile of global BS512 uses G64. One +possible **declared grid**, not a measured cost table, is: + +```json +{ + "physical_k": 3, + "record_confidence_head": true, + "max_samples": 4096, + "cells": [ + {"graph_batch_size": 64, "verifier_budget": 0}, + {"graph_batch_size": 64, "verifier_budget": 96}, + {"graph_batch_size": 64, "verifier_budget": 128}, + {"graph_batch_size": 64, "verifier_budget": 192} + ] +} +``` + +Use `cuda_graph_config.batch_sizes: [64]` and `enable_padding: true` for that +single-G example; the original runtime's token/request limits must retain G64. +If the resolved ladder has more G values, include their native V0 controls and +desired compact cells. Missing or silently clamped coverage is an error. + +The interface can prepare all previously used global batch sizes +128/256/512/1024/1536/2048. With eight attention-DP ranks these map to +G16/G32/G64/G128/G192/G256. With attention DP off, G is the full global batch, +not BS/8. Smaller points are explicit G inputs too; no pilot batch is hardcoded. +Each additional G requires its own captured graph and measured compatible cells. +Preparation does not imply validation of that batch size. Native/GPU validation +is still unperformed for this new tooling, including BS512/1024. + +The JSONL input file must contain at least the largest global cohort size, with +one object per line and a nonempty list of integer `prompt_token_ids`. Use the +actual licensed tokenizer and workload. The driver assigns stable input IDs and +uses the same prompt prefix for repeated cells; prompts and requested output +length must fit `max_seq_len`. + +## Development collection command + +After the complete launch procedure, matching native build and GPU health are +qualified, invoke the CLI under the site's normal MPI/container launcher. For +an already configured, healthy eight-GPU Linux environment the interface is: + +```bash +mpirun -np 8 --bind-to none python -m tensorrt_llm.tools.verify_costs profile \ + --model /path/to/authorized/checkpoint \ + --config /path/to/runtime.yaml \ + --grid /path/to/profile-grid.json \ + --physical-k 3 \ + --inputs /path/to/tokenized-prompts.jsonl \ + --output /path/to/new-private-output-directory \ + --generations 3 --output-tokens 768 --timeout 12180 +``` + +These paths are explicit inputs, not supplied artifacts. `--timeout` is the total +generation-phase request budget after construction/identity, not a guarantee that +startup fits a scheduler lease. Reserve the full startup, model, export and +shutdown budget in the outer launcher. Do not restart on quiet logs or an +observation timeout. The output directory must be new and its parent must exist; +no existing output is overwritten. + +The result contains one `rank-R-record-N.json` per cell/generation/rank. It records +actual replay keys, lease-completed GPU events, request/sampler/KV lifecycle joins, +fixed-length original responses and closure. Completed evidence is retained on +failure with `run_complete: false`; a failure record does not qualify costs. +Successful shutdown is required before raw records can be exported as complete. + +## Export and inspect raw measured costs + +Reduction is stdlib-only and can run without importing native TensorRT-LLM: + +```bash +python tensorrt_llm/tools/verify_costs.py export \ + --raw /path/to/new-private-output-directory/rank-*-record-*.json \ + --output /path/to/new-measured-summary.json \ + --exclude-generations 1 --warmup-samples 2 \ + --minimum-generation-samples 12 --minimum-samples 24 +``` + +The reducer joins exact replay ordinals and executor iterations across **all** +ranks, takes each joined rank maximum, then reports the retained median. It does +not average ranks or sum overlapping spans. It preserves raw-file hashes, +per-rank samples and every exclusion. It rejects duplicate JSON keys, nonfinite +or nonpositive durations, unequal scopes/identities, missing declared cells or +generations, missing/duplicate ranks, reused/reset ordinals, changed observed +layouts and incomplete generation/run/shutdown closure. +Each declared cell must observe one consistent graph key over its generations. +If the workload crosses different sequence-length/attention graph regimes, the +collector refuses to pool them; prepare separate appropriately declared runs. +Each retained generation needs at least 12 exact samples after exclusions, and +each cell needs at least 24 pooled retained samples. A large count in one +generation does not waive another generation's floor. + +Cost keys contain physical K, graph G/V, attention layout, captured-head mode and +the observed graph key. Rows with equal counts but different layouts are not +pooled. Inspect the source/runtime/model/topology identity before comparing or +reusing a summary. Raw summary structure and timing scope remain different from +the runtime's schema-v2 prediction table; renaming a file does not convert it. + +## Fit existing paired confidence observations on CPU + +The standalone fitter accepts tensor-only `.pt` shards containing +`logits: [N, K]`, binary cumulative `prefix_mask: [N, K]` and +`meta.pairing: "draft_seq_ring"`. The labels must be observed against the same +stamped draft sequence and request incarnation as the raw logits, with full +physical-K verification. A shortened compact verification cannot provide labels +for its unverified suffix. This producer contract is not inferred from tensor +dimensions or from a metadata string alone. + +Fit the established 41-point, left-to-right cumulative-survival ECE objective: + +```bash +python tensorrt_llm/tools/dspark_calibration.py \ + --shards /path/to/private/paired-observations-*.pt \ + --physical-k 3 --output /path/to/new-sts.json +``` + +The shell expands the shard glob. The default floor is 1,000 paired rows; use +representative data rather than interpreting a small fixture as a calibrated +head. Physical K is explicit, not fixed to the profiling pilot. The output +retains both `sts_temperatures` and `temperatures`, the original shard hashes, +sample count and before/after ECE. Degenerate prefix positions are reported, +not promoted as informative fits. Inputs use PyTorch's tensor-only safe loader; +duplicate shards, wrong pairing/width, nonfinite data and invalid prefix order +are rejected. Existing output files are never overwritten. + +This command has CPU math/CLI coverage, including exact agreement with the +established fitter. The new public paired-observation collector is still a +separate integration task: the graph-cost collector does not emit these shards. +No private collector command or environment variable is a public prerequisite +advertised here. CPU fitting does not certify GPU freshness, task accuracy or +positive scheduling economics. Ordinary runtime reuse of a valid supplied STS +vector remains available without fitting it again. + +## Next stages and validation boundary + +A measured summary deliberately sets `production_sps_table_eligible`, +`production_admission_qualified`, `policy_overhead_measured` and +`task_accuracy_qualified` to false. Do not rename it to +`confidence_sps_table_path` or create a positive receipt from graph timings alone. + +Confidence scheduling remains default-off. With explicit enablement, the +default reuse path accepts structurally valid supplied costs and calibration, +with clear warnings when freshness and performance economics are unchecked. +It retains physical-K/graph/layout coverage, finite positive values and ordinary +runtime/acceptance safety. Fresh costs/calibration and positive economic +admission are not universal prerequisites to a fair measured comparison. + +The candidate [self-generation workflow](CONFIDENCE_REPRODUCTION.md) includes +whole-step cost export, correctly paired calibration collection/fitting, and +ordinary static/dynamic execution plus comparison. The graph-replay mode keeps +its narrower recorded scope; it does not price online policy, layout preparation +or end-to-end request execution. +Measure those exposed costs separately when evaluating prediction quality or +scheduling headroom. Do not double-count overlapping spans. + +Users choosing `validate_confidence_artifact_freshness: true` additionally need +an actual destination fingerprint and the optional workload-economic checks, +including measured enabled-path overhead. Native decisions pay that overhead +too. A positive receipt must come from genuine measured/replay evidence, never +from graph timings alone. Neither this collector nor the CPU fitter fabricates +one. The public strict-validation producer remains an explicit pending piece. + +A clean-checkout colleague run must ultimately verify requested versus resolved +enablement, genuine compact routing, completed outputs/task checks and a fair +matched static-versus-enabled benchmark. CPU fake-event tests and smoke imports +are not CUDA replay, full-model accuracy, production admission or performance proof. + +The reusable measurement/export module is +`tensorrt_llm/tools/verify_costs.py`; DSpark shape/runtime adapters are under +`_torch/speculative/dspark_cost_profile.py`, +`_torch/speculative/dspark_cost_profiler.py` and the explicit profile-only engine +and request owner under `_torch/pyexecutor`. Future workload adapters can reuse +timing leases, exact-rank reduction and export while supplying their own execution +shape and layout identity. This does not assume unpublished model-specific details. + +Focused developer controls are: + +```bash +python -B tests/unittest/tools/test_dspark_cost_export.py +python -B tests/unittest/tools/test_dspark_calibration.py +python -B tests/unittest/tools/test_dspark_observations.py +python -m pytest -q tests/unittest/executor/test_dspark_reproduction_factory.py +``` + +The artifact and observation suites use CPU data and finite lifecycle controls; +they do not qualify CUDA events or real distributed execution. The factory suite +needs a matching native import environment but does not construct a model. It +covers the ordinary public argument boundary and the explicit observation/profile +constructor seams, including resource-governor rejection for the profile owner. +It also checks the imported parent replay keyword/default contract and preserves +the same position-offset object in native and measured profile replay branches. +It also covers native-collection and enabled-observation graph parents, including +the same position-offset forwarding and routing bookkeeping. +The file contains fourteen parameter cases; its complete matching-runtime pytest run +is distinct from selected-source CPU boundary controls. +These controls are not full-model GPU, real-MPI, accuracy or performance results. diff --git a/examples/models/core/deepseek_v4/README.md b/examples/models/core/deepseek_v4/README.md index ddac8fa83a45..fc956dad5651 100644 --- a/examples/models/core/deepseek_v4/README.md +++ b/examples/models/core/deepseek_v4/README.md @@ -481,6 +481,13 @@ receipt generation and performance qualification are external, and the composed build, GPU/full-model, accuracy and performance validation. Optional mixed scheduling has the separate evidence and route requirements below; no composed-tip GPU qualification is implied. +The source includes candidate [offline cost profiling](COST_PROFILING.md), +[cost export and artifact reuse](COST_EXPORT.md), and +[calibration collection plus matched static/dynamic commands](CONFIDENCE_REPRODUCTION.md). +These use ordinary runtime configuration and configurable physical K and batch size, without +private harnesses or cluster-specific inputs. Their end-to-end validation and publication status +remain as stated in those guides; source availability alone is not GPU qualification. + `K` is the physical draft length, `G` is the rank-local request count for a captured graph, and `V` is its verifier-token budget. In the cost table, `V=0` denotes the native full-width path; positive `V` values denote measured compact layouts. The table must cover the configured physical @@ -570,10 +577,12 @@ With `validate_confidence_artifact_freshness: true`, provide all four prototype - `confidence_admission_receipt_sha256`: the expected lowercase SHA256 of the **receipt file bytes**, obtained from a trusted, separately reviewed deployment package. -There is no supported public cost-profile/calibration/receipt-generation command in this source -change. Self-generation requires the separate profiling tooling and its validated instructions; -the fields above are consumers, not a cost collector. Do not manufacture costs or an admitting -receipt from unit-test fixtures. A strict receipt's producer must preserve the measured cost, +The candidate [self-generation workflow](CONFIDENCE_REPRODUCTION.md) provides separate cost +profiling/export and correctly paired calibration collection/fitting commands. They do not +generate a positive strict-admission receipt. The fields above consume artifacts; strict +freshness validation is optional and remains separate from default artifact reuse. Do not +manufacture costs or an admitting receipt from unit-test fixtures. A strict receipt's producer +must preserve the measured cost, workload, calibration, selector-replay and runtime evidence and review their applicability to the actual deployment. A source, calibration, physical-`K` or workload change requires reassessment; historical scores or timing results cannot be relabeled as new-source qualification. diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 7fcdda9d4be5..bab0d8dbbcba 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -3997,6 +3997,7 @@ def create_py_executor_instance( execution_stream: Optional[torch.cuda.Stream] = None, dwdp_manager: Optional[DwdpManager] = None, max_num_sequences: Optional[int] = None, + executor_type: Optional[type[PyExecutor]] = None, ) -> PyExecutor: set_low_latency_dispatch( getattr(llm_args, 'enable_low_latency_host_dispatch', False)) @@ -4366,7 +4367,8 @@ def create_py_executor_instance( executor_max_seq_len = min(executor_max_seq_len, int(decoder_position_limit)) - return PyExecutor( + executor_type = PyExecutor if executor_type is None else executor_type + return executor_type( resource_manager, scheduler, model_engine=model_engine, diff --git a/tensorrt_llm/_torch/pyexecutor/dspark_observation_runtime.py b/tensorrt_llm/_torch/pyexecutor/dspark_observation_runtime.py new file mode 100644 index 000000000000..d942ff8cf789 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/dspark_observation_runtime.py @@ -0,0 +1,341 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Explicit offline observation factory over normal native/automatic execution. + +No fixed-cell/profile-only engine is used here. Calibration alone evaluates the +confidence head under native full-K verification and copies observations before +the SAME normal sampler event; benchmark arms add no calibration copies. +""" + +from __future__ import annotations + +from functools import partial +from typing import Any + +from tensorrt_llm.tools.dspark_observations import NativePairs, ProposalSnapshot + + +class ObservationExecutorMixin: + """Thin lifecycle joins; every forward, sampler/update and error delegates.""" + + def __init__(self, *args, **kwargs): + self._observation_inputs = {} + self._observation_states = {} + self._observation_pending = {} + self._observation_forward = None + self._observation_failures = [] + super().__init__(*args, **kwargs) + self._native_pairs = NativePairs(self.model_engine.spec_config.max_draft_len) + + def _handle_errors( + self, + error_msg=None, + *, + requests=None, + charge_budget=True, + fatal_is_collective_aligned=False, + ): + self._observation_failures.append(str(error_msg or "Original executor error")) + return super()._handle_errors( + error_msg, + requests=requests, + charge_budget=charge_budget, + fatal_is_collective_aligned=fatal_is_collective_aligned, + ) + + def _forward_step(self, batch, new_tensors_device=None): + if self.is_warmup or not self.model_engine._collect_calibration: + return super()._forward_step(batch, new_tensors_device) + if self._observation_forward is not None: + raise RuntimeError("Overlapping unpaired observation forwards") + requests = tuple(r for r in batch.generation_requests if not r.is_dummy) + source = self._observation_inputs.get(id(new_tensors_device)) + input_host = None + if source is not None: + if source.state.device is not new_tensors_device: + raise RuntimeError("Original device input incarnation changed") + # Issued before original forward can reuse persistent output storage. + # The next original sample_async event covers this same-stream copy. + input_host = self.sampler._copy_to_host(new_tensors_device.next_draft_tokens) + elif requests: + self._native_pairs.stats["bootstrap_or_no_bound_input"] += len(requests) + # Existing source stamp scalar is observation identity only. It does + # not alter proposal/acceptance math or construct a scheduling policy. + draft_sequence = self.model_engine.model.spec_worker.bump_draft_seq() + 1 + self._observation_forward = ( + self.iter_counter, + requests, + source, + input_host, + draft_sequence, + ) + result = super()._forward_step(batch, new_tensors_device) + if result is None: + raise RuntimeError("Original forward failed; observations cannot be promoted") + return result + + def _sample_async(self, batch, output): + if self.is_warmup or not self.model_engine._collect_calibration: + return super()._sample_async(batch, output) + worker = self.model_engine.model.spec_worker + confidence = worker.staged_confidence_buffer() + if confidence is None or self._observation_forward is None: + raise RuntimeError("Offline full-K collector has no produced confidence/forward") + iteration, requests, source, input_host, draft_sequence = self._observation_forward + rows = {} + for request in requests: + row = worker.confidence_row_for(request.py_request_id) + if row is None or request.py_seq_slot is None: + raise RuntimeError("Original request has no worker confidence/sampler slot") + rows[request.py_request_id] = (request, int(request.py_seq_slot), int(row)) + # Already-produced full-K logits are copied with the original sampler's + # helper BEFORE it records its normal event. No extra event or wait. + host = self.sampler._copy_to_host(confidence) + stamps = self.sampler._copy_to_host(worker.confidence_stamp_buffer()) + state = super()._sample_async(batch, output) + if state is None or id(state) in self._observation_pending: + raise RuntimeError("Missing or reused original sampler state") + snapshot = ProposalSnapshot(iteration, state, host, rows, stamps, draft_sequence) + self._observation_inputs[id(state.device)] = snapshot + self._observation_states[id(state)] = snapshot + self._observation_pending[id(state)] = (iteration, requests, source, input_host) + self._observation_forward = None + return state + + def _update_requests(self, state, resource_manager=None): + if self.is_warmup or not self.model_engine._collect_calibration: + return super()._update_requests(state, resource_manager) + pending = self._observation_pending.get(id(state)) + snapshot = self._observation_states.get(id(state)) + if pending is None or snapshot is None or snapshot.state is not state: + raise RuntimeError("Unknown or consumed observation sampler state") + before = {id(request): request.py_decoding_iter for request in state.requests} + result = super()._update_requests(state, resource_manager) + # Original SpecSampler.update_requests performed its existing synchronize + # and updates. Parent errors are fatal to publication, not waived here. + if self._observation_failures: + raise RuntimeError("Original sampler/update error; no paired-data publication") + snapshot.updated = True + iteration, requests, source, input_host = pending + by_object = {id(request): index for index, request in enumerate(state.requests)} + for request in requests: + index = by_object.get(id(request)) + if index is None: + raise RuntimeError("Original generation request absent from sampler state") + if request.py_decoding_iter <= before[id(request)]: + self._native_pairs.stats["retired_or_not_advanced"] += 1 + continue + if source is None: + continue + if request.py_request_id not in source.requests: + # A finished-context output has only placeholder draft rows, + # not a confidence-scored block from a generation forward. + self._native_pairs.stats["prefill_placeholder_no_scored_proposal"] += 1 + continue + slot = int(request.py_seq_slot) + token_window = ( + int(state.host.verify_lens[index]) + if state.host.verify_lens is not None + else state.runtime_draft_len + 1 + ) + verified = min(int(state.draft_lens[index]), token_window - 1) + if verified != self._native_pairs.physical_k: + # First-draft bootstrap has no scored full-K proposal. Never + # reinterpret its placeholder confidence as calibration data. + self._native_pairs.stats["bootstrap_or_partial"] += 1 + continue + self._native_pairs.record( + source=source, + request=request, + target_iteration=iteration, + input_proposal=input_host[slot], + target_slot=slot, + accepted=int(state.host.new_tokens_lens[slot]) - 1, + verified=verified, + ) + del self._observation_pending[id(state)] + return result + + def _observation_drain_predicates(self): + previous = self.previous_batch + state = previous.sample_state if previous is not None else None + snapshot = self._observation_states.get(id(state)) if state is not None else None + return { + "failures": len(self._observation_failures), + "active_requests": len(self.active_requests), + "waiting_requests": len(self.waiting_queue), + "pending_samples": len(self._observation_pending), + "forward_pending": self._observation_forward is not None, + "previous_batch_present": previous is not None, + "pending_matches_previous": ( + state is not None + and len(self._observation_pending) == 1 + and id(state) in self._observation_pending + and snapshot is not None + and snapshot.state is state + and not snapshot.updated + ), + } + + def observation_receipt(self, *, reset=False, control=None): + with self.control_action(drain=True): + before = self._observation_drain_predicates() + consumed = False + if self.model_engine._collect_calibration: + if control is None: + raise RuntimeError("Offline drain needs the original gang-control communicator") + peers = control.allgather(before) + quiescent = all( + row["failures"] == 0 + and row["active_requests"] == 0 + and row["waiting_requests"] == 0 + and not row["forward_pending"] + for row in peers + ) + no_tail = all(row["pending_samples"] == 0 for row in peers) + exact_tail = all( + row["pending_samples"] == 1 and row["pending_matches_previous"] for row in peers + ) + if not quiescent or not (no_tail or exact_tail): + raise RuntimeError( + "Original requests/updates did not naturally drain: " + f"gang_predicates={peers!r}" + ) + if exact_tail: + # Every rank is paused at the original exclusive boundary. + # Complete the same existing sampler/update/resource tail + # before checking the observer's stronger publication gate. + self._consume_previous_batch_for_rebalance() + consumed = True + after = self._observation_drain_predicates() + if ( + self._observation_failures + or self.active_requests + or len(self.waiting_queue) + or self._observation_pending + or self._observation_forward is not None + ): + raise RuntimeError( + "Original requests/updates did not naturally drain: " + f"before={before!r}, after={after!r}" + ) + runner = self.model_engine.cuda_graph_runner + result = { + "drain": { + "before": before, + "after": after, + "original_previous_batch_consumed": consumed, + }, + "routing": dict(runner.observed_routes), + "graph_keys": list(runner.observed_keys.values()), + "resolved_confidence_enabled": bool( + self.model_engine.spec_config.enable_confidence_scheduling + ), + "runtime_confidence_enabled": bool( + getattr(self.model_engine, "_dspark_confidence_enabled", False) + ), + } + if reset: + runner.observed_routes = {"native_replays": 0, "compact_replays": 0} + runner.observed_keys.clear() + return result + + +def create_observation_executor( + *, llm_args, collect_calibration=False, checkpoint_dir=None, tokenizer=None +): + """Only this public offline caller gets observation hooks; normal serving is unchanged.""" + from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( + CUDAGraphRunner, + DSparkCUDAGraphRunner, + KeyType, + ) + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner + from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + from tensorrt_llm._torch.pyexecutor.py_executor_creator import ( + _RuntimeConstructors, + create_py_executor, + ) + from tensorrt_llm._torch.speculative.dspark import DSv4DSparkWorker, resolve_confidence_head + from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig + + spec = llm_args.speculative_config + if not isinstance(spec, DSparkDecodingConfig) or not spec.draft_is_embedded_in_target: + raise ValueError("The DSpark adapter requires embedded chain drafting") + if collect_calibration and (spec.enable_confidence_scheduling or spec.draft_len_schedule): + raise ValueError( + "Offline calibration requires native full-K, not policy-trimmed verification" + ) + + # Retain PR5's ordinary native-OFF versus DSpark-ON graph implementation. + graph_base = DSparkCUDAGraphRunner if spec.enable_confidence_scheduling else CUDAGraphRunner + + class ObservedGraphs(graph_base): + def __init__(self, *args, **kwargs): + self.observed_routes = {"native_replays": 0, "compact_replays": 0} + self.observed_keys = {} + super().__init__(*args, **kwargs) + + def replay(self, key, inputs, *, position_id_offsets=None): + result = super().replay(key, inputs, position_id_offsets=position_id_offsets) + key = KeyType(*key) # same accepted tuple/KeyType contract as the original runner + if key.is_first_draft or key.draft_len != spec.max_draft_len: + return result + raw = dict(key._asdict()) + # Existing KeyType carries a SampleType enum; no arbitrary objects + # or environment values are recorded by this finite graph witness. + raw["sample_type"] = str(raw["sample_type"]) + budget = getattr(key, "ragged_verify_bucket", None) + compact = budget is not None and 0 < budget < key.batch_size * (spec.max_draft_len + 1) + self.observed_routes["compact_replays" if compact else "native_replays"] += 1 + self.observed_keys[repr(key)] = raw + return result + + class ObservationDecoderRunner(DecoderRunner): + def _initialize_cuda_graph_runner(self): + if collect_calibration: + worker = self.model.spec_worker + if type(worker) is not DSv4DSparkWorker or worker._win_inited: + raise RuntimeError( + "Calibrated head must be enabled before original graph/storage setup" + ) + if resolve_confidence_head(self.model.draft_model) is None: + raise ValueError("Actual drafter has no confidence head") + worker.return_confidence = True # head evaluation only, no planner/selection + return super()._initialize_cuda_graph_runner() + + def _create_cuda_graph_runner(self, config): + return ObservedGraphs(config) + + class ObservationEngine(PyTorchModelEngine): + def __init__(self, **kwargs: Any): + self._collect_calibration = collect_calibration + super().__init__(**kwargs) + + def _initialize_decoder_runner(self, runner_cls, runner_config=None, **runner_kwargs): + if runner_cls is not DecoderRunner: + raise ValueError("Offline DSpark observation requires the original decoder runner") + return super()._initialize_decoder_runner( + ObservationDecoderRunner, runner_config, **runner_kwargs + ) + + class ObservationExecutor(ObservationExecutorMixin, PyExecutor): + pass + + constructors = _RuntimeConstructors( + model_engine=ObservationEngine, + executor=ObservationExecutor, + kv_cache_creator=KvCacheCreator, + ) + return create_py_executor( + llm_args=llm_args, + checkpoint_dir=checkpoint_dir, + tokenizer=tokenizer, + runtime_constructors=constructors, + ) + + +def observation_factory(*, collect_calibration=False): + return partial(create_observation_executor, collect_calibration=collect_calibration) diff --git a/tensorrt_llm/_torch/pyexecutor/dspark_profile_engine.py b/tensorrt_llm/_torch/pyexecutor/dspark_profile_engine.py new file mode 100644 index 000000000000..ecf50df546cf --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/dspark_profile_engine.py @@ -0,0 +1,462 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Non-serving fixed-cell DSpark profiling on the original model engine.""" + +from __future__ import annotations + +import hashlib +import os +from collections.abc import Sequence +from typing import Any + +import torch + +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( + CUDAGraphRunner, + CUDAGraphRunnerConfig, + DSparkCUDAGraphRunner, +) +from tensorrt_llm._torch.pyexecutor.dspark_profile_graph_storage import ( + snapshot_metadata_storage, + validate_metadata_storage, +) +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner +from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager +from tensorrt_llm._torch.pyexecutor.sampler import SampleStateTensors +from tensorrt_llm._torch.pyexecutor.sampler.sampler_common import SampleType +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.speculative.dspark import DSv4DSparkWorker, resolve_confidence_head +from tensorrt_llm._torch.speculative.dspark_cost_profile import ( + DSparkCostProfileConfig, + DSparkProfileCell, +) +from tensorrt_llm._torch.speculative.spec_sampler_base import SampleStateTensorsSpec +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig, PrefillCudaGraphBackend +from tensorrt_llm.tools.verify_costs import EventLeasePool + +__all__ = ["DSparkCostProfileEngine"] + + +def _require_shape_agreement(rows: Sequence[Sequence[int]]) -> None: + if not rows or any(len(row) != 5 or row[0] != 1 for row in rows): + raise ValueError("Every profile rank must provide a valid decode-only cell") + if any(tuple(row[1:]) != tuple(rows[0][1:]) for row in rows): + raise ValueError("Profile ranks disagree on G, physical K, V or replay ordinal") + + +def _profile_contract_row(profile: DSparkCostProfileConfig, valid_ladder: bool) -> list[int]: + digest = hashlib.sha256(profile.model_dump_json().encode()).digest() + return [ + int(valid_ladder), + profile.physical_k, + int(profile.record_confidence_head), + profile.max_samples, + *(int.from_bytes(digest[index : index + 8], signed=True) for index in range(0, 32, 8)), + ] + + +def _require_profile_agreement(rows: Sequence[Sequence[int]]) -> None: + if ( + not rows + or any(len(row) != 8 or row[0] != 1 for row in rows) + or any(tuple(row[1:]) != tuple(rows[0][1:]) for row in rows) + ): + raise ValueError( + "Every profile rank must agree on the construction-only grid and head mode" + ) + + +def _valid_previous_state(state: Any, requests: Sequence[Any], physical_k: int) -> bool: + """Validate original overlap tensor extents without reading device values.""" + if state is None: + return all(len(request.py_draft_tokens) == physical_k for request in requests) + if not isinstance(state, SampleStateTensorsSpec): + return False + tensors = (state.new_tokens, state.new_tokens_lens, state.next_draft_tokens) + if any( + not isinstance(value, torch.Tensor) + or not value.is_cuda + or value.dtype not in (torch.int32, torch.int64) + or not value.is_contiguous() + for value in tensors + ): + return False + tokens, lengths, drafts = tensors + if ( + tokens.ndim != 3 + or tokens.shape[0] != physical_k + 1 + or tokens.shape[2] != 1 + or lengths.ndim != 1 + or drafts.ndim != 2 + or drafts.shape[1] != physical_k + or tokens.shape[1] != lengths.shape[0] + or drafts.shape[0] != lengths.shape[0] + or len({value.device for value in tensors}) != 1 + ): + return False + slots = [getattr(request, "py_batch_idx", None) for request in requests] + return all(type(slot) is int and 0 <= slot < lengths.shape[0] for slot in slots) and len( + set(slots) + ) == len(slots) + + +class _DSparkProfileWorker(DSv4DSparkWorker): + def __init__( + self, + spec_config: DSparkDecodingConfig, + mapping: Any, + *, + record_head: bool, + use_separate_draft_kv_cache: bool, + ) -> None: + super().__init__(spec_config, mapping, use_separate_draft_kv_cache) + self.return_confidence = record_head + + def _build_verify_planner(self, draft_model: torch.nn.Module, block_size: int) -> None: + if block_size != self.spec_config.max_draft_len: + raise ValueError("Profile head width must match the physical draft block") + if resolve_confidence_head(draft_model) is None: + raise ValueError("Head-ON profiling requires the loaded confidence head") + # Buffers and the original head forward are retained; no serving planner is constructed. + self.verify_planner = None + + +class _DSparkProfileGraphs(DSparkCUDAGraphRunner): + _snapshot_graph_tensor_ptrs = staticmethod(snapshot_metadata_storage) + _assert_graph_tensor_ptrs_stable = staticmethod(validate_metadata_storage) + + def __init__(self, config: CUDAGraphRunnerConfig, profile: DSparkCostProfileConfig) -> None: + super().__init__(config) + self._profile = profile + self._dspark_trims_submitted_tokens = True + self.active_profile_cell: DSparkProfileCell | None = None + self.active_request_ids: tuple[int, ...] = () + self._profile_pool = EventLeasePool(lambda: torch.cuda.Event(enable_timing=True)) + self._profile_pool.initialize() + torch.cuda.current_stream().synchronize() + self._profile_records: list[dict[str, Any]] = [] + self._profile_ptrs: dict[Any, tuple[dict[str, Any], dict[str, Any]]] = {} + + def capture(self, key: Any, *args: Any, **kwargs: Any) -> Any: + output = super().capture(key, *args, **kwargs) + metadata = self.graph_metadata[key] + self._profile_ptrs[key] = ( + self._snapshot_graph_tensor_ptrs(metadata["attn_metadata"]), + self._snapshot_graph_tensor_ptrs(metadata["spec_metadata"]) + if metadata["spec_metadata"] is not None + else {}, + ) + return output + + def maybe_get_cuda_graph(self, batch: ScheduledRequests, *args: Any, **kwargs: Any) -> Any: + result = super().maybe_get_cuda_graph(batch, *args, **kwargs) + cell = self.active_profile_cell + if cell is not None: + key = result[2] + if key is None or key not in self.graphs: + raise RuntimeError("A real profile cell requires its previously captured graph") + expected_v = None if cell.verifier_budget == 0 else cell.verifier_budget + if (key.batch_size, key.draft_len, key.ragged_verify_bucket) != ( + cell.graph_batch_size, + self._profile.physical_k, + expected_v, + ): + raise RuntimeError("Observed replay key does not match the declared profile cell") + if key.num_contexts != 0 or not key.is_all_greedy_sample: + raise RuntimeError("The fixed-cell pilot measures greedy decode-only graphs") + return result + + def replay( + self, + key: Any, + current_inputs: dict[str, Any], + *, + position_id_offsets: torch.Tensor | None = None, + ) -> Any: + cell = self.active_profile_cell + if cell is None: + return super().replay(key, current_inputs, position_id_offsets=position_id_offsets) + if len(self._profile_records) >= self._profile.max_samples: + raise RuntimeError("Preallocated profile timing capacity exhausted") + metadata = self.graph_metadata[key] + attn_ptrs, spec_ptrs = self._profile_ptrs[key] + storage_context = { + "rank": self.config.mapping.rank, + "graph_key": tuple( + value.name if isinstance(value, SampleType) else value for value in key + ), + } + self._assert_graph_tensor_ptrs_stable( + metadata["attn_metadata"], attn_ptrs, context=storage_context + ) + if metadata["spec_metadata"] is not None: + self._assert_graph_tensor_ptrs_stable( + metadata["spec_metadata"], spec_ptrs, context=storage_context + ) + lease = self._profile_pool.begin() + output = super().replay(key, current_inputs, position_id_offsets=position_id_offsets) + self._profile_pool.end(lease) + self._profile_records.append( + { + "sample_index": len(self._profile_records), + "_event_lease": lease, + "rank": self.config.mapping.rank, + "graph_batch_size": cell.graph_batch_size, + "physical_k": self._profile.physical_k, + "verifier_budget": cell.verifier_budget, + "request_ids": self.active_request_ids, + "graph_key": tuple( + value.name if isinstance(value, SampleType) else value for value in key + ), + } + ) + return output + + def complete_sample(self, index: int) -> None: + """Release events after the original sampler/update completion boundary.""" + record = self._profile_records[index] + lease = record["_event_lease"] + elapsed = self._profile_pool.complete(lease) + del record["_event_lease"] + record["replay_entry_elapsed_ms"] = elapsed + + def completed_profile_samples(self) -> list[dict[str, Any]]: + """Return completed raw measurements without waiting or inventing costs.""" + if self._profile_pool.pending: + raise RuntimeError("Profile generation still owns unfinished event leases") + if any("replay_entry_elapsed_ms" not in row for row in self._profile_records): + raise RuntimeError("Profile replay has no original update completion") + return [dict(row) for row in self._profile_records] + + +class _DSparkProfileDecoderRunner(DecoderRunner): + """Offline-only capture/state owner, using the current decoder lifecycle.""" + + def __init__(self, model, config, *, profile, **runner_kwargs): + self._cost_profile = profile + super().__init__(model, config, **runner_kwargs) + + @property + def _ragged_verification_enabled(self) -> bool: + return True + + def _initialize_cuda_graph_runner(self): + old_worker = self.model.spec_worker + if type(old_worker) is not DSv4DSparkWorker or old_worker._win_inited: + raise RuntimeError("Profile worker must be installed before its first initialization") + if old_worker.staged_confidence_buffer() is not None or self.attn_metadata is not None: + raise RuntimeError( + "Profile construction cannot replace initialized state or graph addresses" + ) + matches = [index for index, entry in enumerate(self.model.epilogue) if entry is old_worker] + if len(matches) != 1: + raise RuntimeError("Original speculative worker must have one epilogue owner") + worker = _DSparkProfileWorker( + self._config.spec_config, + self.mapping, + record_head=self._cost_profile.record_confidence_head, + use_separate_draft_kv_cache=self.model.use_separate_draft_kv_cache, + ) + worker.set_draft_model(self.model.draft_model) + self.model.spec_worker = worker + self.model.epilogue[matches[0]] = worker + self._dspark_trims_submitted_tokens = True + self.previous_verify_lens = [] + return super()._initialize_cuda_graph_runner() + + def _create_cuda_graph_runner(self, config: CUDAGraphRunnerConfig) -> CUDAGraphRunner: + return _DSparkProfileGraphs(config, self._cost_profile) + + def _get_graphs_to_capture(self, cuda_graph_batch_sizes: list[int]) -> list[tuple[int, ...]]: + return self._cost_profile.graph_shapes(cuda_graph_batch_sizes) + + def ragged_verify_token_buckets(self, padded_bs: int) -> list[int]: + return sorted( + cell.verifier_budget + for cell in self._cost_profile.cells + if cell.graph_batch_size == padded_bs and cell.verifier_budget > 0 + ) + + +class DSparkCostProfileEngine(PyTorchModelEngine): + """Explicit offline engine; the ordinary public serving factories cannot create it. + + Args: + profile: Validated construction-only cell grid. + engine_kwargs: Original PyTorchModelEngine construction arguments. The + DSpark serving configuration must keep confidence scheduling off. + """ + + def __init__(self, *, profile: DSparkCostProfileConfig, **engine_kwargs: Any) -> None: + self._cost_profile = profile + self._profile_call_active = False + spec = engine_kwargs.get("spec_config") + if not isinstance(spec, DSparkDecodingConfig): + raise TypeError("Profile-only construction requires DSparkDecodingConfig") + if spec.max_draft_len != profile.physical_k or not spec.draft_is_embedded_in_target: + raise ValueError("Profile-only construction requires embedded DSpark with matching K") + if ( + spec.enable_confidence_scheduling + or spec.enable_fused_confidence_scheduler + or getattr(spec, "confidence_mixed_evidence", None) is not None + or spec.draft_len_schedule is not None + ): + raise ValueError("Profile-only construction cannot enable serving confidence policy") + if any( + getattr(spec, name, None) + for name in ( + "confidence_sps_table_path", + "confidence_sps_live_fingerprint_path", + "confidence_admission_receipt_path", + "confidence_admission_receipt_sha256", + "confidence_sts_path", + ) + ): + raise ValueError( + "Profile-only construction does not consume serving admission artifacts" + ) + if os.environ.get("TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS"): + raise ValueError("Fixed-cell costs require real acceptance, not forced acceptance") + super().__init__(**engine_kwargs) + + def _initialize_decoder_runner(self, runner_cls, runner_config=None, **runner_kwargs): + valid_ladder = True + try: + self._cost_profile.validate_graph_ladder(self._cuda_graph_batch_sizes) + except ValueError: + valid_ladder = False + if self.mapping.tp_size > 1 and self.dist is None: + raise ValueError("Multi-rank profiling requires the original distributed communicator") + row = _profile_contract_row(self._cost_profile, valid_ladder) + rows = self.dist.tp_allgather_int64(row).tolist() if self.mapping.tp_size > 1 else [row] + _require_profile_agreement(rows) + if ( + runner_cls is not DecoderRunner + or self.is_multimodal + or self._is_encoder_decoder_model() + or self.max_beam_width != 1 + or self.mapping.pp_size != 1 + or self.mapping.cp_size != 1 + or self.cuda_graph_config is None + or not self._cuda_graph_padding_enabled + or self.prefill_cuda_graph_backend != PrefillCudaGraphBackend.DISABLED + or self.llm_args.lora_config is not None + or self.llm_args.enable_in_graph_sampling + ): + raise ValueError("Unsupported profile-only model, topology or graph configuration") + return super()._initialize_decoder_runner( + _DSparkProfileDecoderRunner, + runner_config, + profile=self._cost_profile, + **runner_kwargs, + ) + + def _set_warmup_ragged_windows(self, batch, verify_bucket, draft_len): + if not isinstance(self._runner, _DSparkProfileDecoderRunner): + raise RuntimeError("Prescribed windows require the offline decoder owner") + return self._runner._set_warmup_ragged_windows(batch, verify_bucket, draft_len) + + def forward(self, *args: Any, **kwargs: Any) -> Any: + raise RuntimeError("Profile-only engines cannot execute the ordinary serving forward API") + + def forward_profile_prefill( + self, batch: ScheduledRequests, resource_manager: ResourceManager + ) -> Any: + """Initialize genuine request KV using the original unmeasured context forward.""" + valid = ( + batch.num_context_requests > 0 + and not batch.num_generation_requests + and not batch.num_encoder_requests + and not self._profile_call_active + ) + row = [int(valid), -1, self._cost_profile.physical_k, 0, -1] + rows = self.dist.tp_allgather_int64(row).tolist() if self.mapping.tp_size > 1 else [row] + _require_shape_agreement(rows) + return super().forward(batch, resource_manager=resource_manager) + + def forward_profile_cell( + self, + cell: DSparkProfileCell, + batch: ScheduledRequests, + resource_manager: ResourceManager, + *, + new_tensors_device: SampleStateTensors | None = None, + cache_indirection_buffer: torch.Tensor | None = None, + ) -> Any: + """Borrow real scheduled state and run one declared cell through the original forward. + + The caller retains ownership of request/sampler/KV updates and must + advance them using the original executor lifecycle before the next call. + """ + requests = batch.generation_requests + local_valid = ( + cell in self._cost_profile.cells + and not self._profile_call_active + and batch.num_context_requests == 0 + and batch.num_encoder_requests == 0 + and len(requests) == cell.graph_batch_size + and len({request.py_request_id for request in requests}) == len(requests) + and all( + not request.is_dummy + and request.py_beam_width == 1 + and not request.py_is_first_draft + and request.py_seq_slot is not None + for request in requests + ) + and self.runtime_draft_len == self._cost_profile.physical_k + and _valid_previous_state(new_tensors_device, requests, self._cost_profile.physical_k) + ) + row = [ + int(local_valid), + cell.graph_batch_size, + self._cost_profile.physical_k, + cell.verifier_budget, + len(self.cuda_graph_runner._profile_records), + ] + rows = self.dist.tp_allgather_int64(row).tolist() if self.mapping.tp_size > 1 else [row] + _require_shape_agreement(rows) + graph_runner = self.cuda_graph_runner + graph_runner.active_profile_cell = cell + graph_runner.active_request_ids = tuple(request.py_request_id for request in requests) + self._profile_call_active = True + try: + if cell.verifier_budget: + self._set_warmup_ragged_windows( + batch, cell.verifier_budget, self._cost_profile.physical_k + ) + else: + for request in requests: + request.py_verify_len = None + graph_runner.agreed_ragged_bucket = None + before = len(graph_runner._profile_records) + output = super().forward( + batch, + resource_manager=resource_manager, + new_tensors_device=new_tensors_device, + cache_indirection_buffer=cache_indirection_buffer, + ) + if len(graph_runner._profile_records) != before + 1: + raise RuntimeError( + "Profile forward did not execute exactly one original graph replay" + ) + return output + finally: + graph_runner.active_profile_cell = None + graph_runner.active_request_ids = () + self._profile_call_active = False + + def profile_observations(self) -> dict[str, Any]: + """Return raw observations, never an admitted table or enabled-performance result.""" + return { + "schema": "dspark-cost-profile-observation-v1", + "profile_only": True, + "physical_k": self._cost_profile.physical_k, + "record_confidence_head": self._cost_profile.record_confidence_head, + "timing_scope": "original-graph-replay-entry-including-static-copies", + "samples": self.cuda_graph_runner.completed_profile_samples(), + "full_policy_h": None, + "production_sps_table_eligible": False, + "production_admission_qualified": False, + } diff --git a/tensorrt_llm/_torch/pyexecutor/dspark_profile_graph_storage.py b/tensorrt_llm/_torch/pyexecutor/dspark_profile_graph_storage.py new file mode 100644 index 000000000000..607196ff9800 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/dspark_profile_graph_storage.py @@ -0,0 +1,188 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Offline-profiler metadata storage checks; no model or graph operations.""" + +from __future__ import annotations + +import json +from typing import Any + +import torch + +_V4_METADATA = ( + "tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.metadata", + "DeepseekV4TrtllmAttentionMetadata", +) +_TRT_METADATA_TYPES = { + _V4_METADATA, + ("tensorrt_llm._torch.attention.backends.sparse.dsa.metadata", "DSAtrtllmAttentionMetadata"), + ("tensorrt_llm._torch.attention.backends.trtllm", "TrtllmAttentionMetadata"), +} +_V4_BUFFER_MAPS = ( + "new_comp_kv_lens_cuda", + "cu_new_comp_kv_cuda", + "compressed_kv_lens_cuda", + "past_kv_lens_cuda", + "compressed_position_ids_cuda", + "compressed_mask_cuda", + "compress_block_tables", + "sparse_mla_topk_lens", +) + + +def _metadata_type(obj: Any) -> tuple[str, str]: + return type(obj).__module__, type(obj).__name__ + + +def _signature(value: torch.Tensor) -> dict[str, Any]: + return { + "address": value.data_ptr(), + "storage_offset": value.storage_offset(), + "shape": tuple(value.shape), + "stride": tuple(value.stride()), + "dtype": str(value.dtype), + "device": str(value.device), + } + + +def _role(obj: Any, name: str) -> str: + metadata_type = _metadata_type(obj) + if ( + name == "gen_indexer_kv_lens_cuda_runtime" + and metadata_type == _V4_METADATA + and obj.is_cuda_graph + and obj.num_contexts == 0 + and obj.num_generations > 0 + and obj._indexer_compress_ratio == 4 + ): + return "graph_produced_compressed_lengths" + if name == "workspace" and metadata_type in _TRT_METADATA_TYPES and obj.is_cuda_graph: + return "eager_only_workspace" + return "persistent" + + +def _refuse( + obj: Any, name: str, reason: str, before: Any, after: Any, context: dict[str, Any] | None = None +) -> None: + context = {} if context is None else context + details = { + "reason": reason, + "field": name, + "metadata_type": _metadata_type(obj), + "expected": before, + "actual": after, + "rank": context.get("rank"), + "graph_key": context.get("graph_key"), + } + raise RuntimeError( + "Profile graph storage contract failed: " + json.dumps(details, sort_keys=True) + ) + + +def _cuda_tensors(obj: Any, *, context: dict[str, Any] | None = None) -> dict[str, torch.Tensor]: + tensors = { + name: value + for name, value in vars(obj).items() + if isinstance(value, torch.Tensor) and value.is_cuda + } + if _metadata_type(obj) == _V4_METADATA: + for name in _V4_BUFFER_MAPS: + values = getattr(obj, name, None) + if not isinstance(values, dict) or any(type(key) is not int for key in values): + _refuse( + obj, + name, + "Invalid or missing required profile storage map", + "dict with integer keys", + type(values).__name__, + context, + ) + for key, value in values.items(): + if not isinstance(value, torch.Tensor) or not value.is_cuda: + _refuse( + obj, + f"{name}[{key}]", + "Invalid profile CUDA storage", + "CUDA tensor", + type(value).__name__, + context, + ) + tensors[f"{name}[{key}]"] = value + return tensors + + +def _validate_intermediate(obj: Any, value: torch.Tensor) -> None: + if ( + tuple(value.shape) != (obj.num_generations,) + or not value.is_contiguous() + or value.storage_offset() != 0 + or value.dtype != obj.kv_lens_cuda.dtype + or value.device != obj.kv_lens_cuda.device + ): + raise RuntimeError("Invalid profile graph-produced compressed-length tensor") + + +def snapshot_metadata_storage(obj: Any) -> dict[str, Any]: + """Record top-level CUDA tensors and audited V4 persistent buffer maps. + + This is a finite profiler contract, not an exhaustive traversal of arbitrary + runtime objects. Unknown tensors are persistent; no user-provided exemption + list or baseline reset is accepted. + """ + records = {} + for name, value in _cuda_tensors(obj).items(): + role = _role(obj, name) + if role == "graph_produced_compressed_lengths": + _validate_intermediate(obj, value) + records[name] = {"role": role, "signature": _signature(value)} + return {"metadata_type": _metadata_type(obj), "fields": records} + + +def validate_metadata_storage( + obj: Any, expected: dict[str, Any], *, context: dict[str, Any] | None = None +) -> None: + """Refuse genuine storage/layout drift with finite, field-specific details.""" + + def refuse(name: str, reason: str, before: Any, after: Any) -> None: + _refuse(obj, name, reason, before, after, context) + + if _metadata_type(obj) != tuple(expected["metadata_type"]): + refuse( + "", + "metadata type changed", + expected["metadata_type"], + _metadata_type(obj), + ) + actual = _cuda_tensors(obj, context=context) + before_fields = expected["fields"] + if set(actual) != set(before_fields): + changed = sorted(set(actual) ^ set(before_fields))[0] + refuse( + changed, + "CUDA field missing or introduced", + before_fields.get(changed), + _signature(actual[changed]) if changed in actual else None, + ) + for name, value in actual.items(): + before = before_fields[name] + role = _role(obj, name) + after = _signature(value) + if role != before["role"]: + refuse(name, "field role changed", before, {"role": role, "signature": after}) + if role == "graph_produced_compressed_lengths": + try: + _validate_intermediate(obj, value) + except RuntimeError as error: + refuse(name, str(error), before, after) + comparable_before = { + key: item for key, item in before["signature"].items() if key != "address" + } + comparable_after = {key: item for key, item in after.items() if key != "address"} + elif role == "eager_only_workspace": + # The graph uses cuda_graph_workspace, which remains persistent. + comparable_before = {key: before["signature"][key] for key in ("dtype", "device")} + comparable_after = {key: after[key] for key in ("dtype", "device")} + else: + comparable_before, comparable_after = before["signature"], after + if comparable_after != comparable_before: + refuse(name, "storage or layout changed", before, {"role": role, "signature": after}) diff --git a/tensorrt_llm/_torch/pyexecutor/dspark_request_profile.py b/tensorrt_llm/_torch/pyexecutor/dspark_request_profile.py new file mode 100644 index 000000000000..0784e84c8574 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/dspark_request_profile.py @@ -0,0 +1,535 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Explicit offline request ownership over the unchanged executor lifecycle.""" + +from __future__ import annotations + +import hashlib +import json +import math +from contextlib import contextmanager +from dataclasses import dataclass +from functools import partial +from typing import TYPE_CHECKING, Any, Optional + +from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator +from tensorrt_llm._torch.pyexecutor.dspark_profile_engine import ( + DSparkCostProfileEngine, + _valid_previous_state, +) +from tensorrt_llm._torch.pyexecutor.llm_request import ExecutorRequest, LlmRequestState +from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm._torch.pyexecutor.py_executor_creator import ( + _RuntimeConstructors, + create_py_executor, +) + +if TYPE_CHECKING: + from ray.actor import ActorHandle + +__all__ = ["OwnedDSparkProfileExecutor", "create_profile_executor"] + + +def _token_snapshot(request: Any) -> dict[str, Any]: + tokens = list(request.get_tokens(0)) + prompt = tokens[: request.py_orig_prompt_len] + return { + "request_id": request.py_request_id, + "process_local_request_object": id(request), + "seq_slot": request.py_seq_slot, + "batch_idx": request.py_batch_idx, + "decoding_iter": request.py_decoding_iter, + "tokens": tokens, + "original_input_sha256": hashlib.sha256( + json.dumps(prompt, separators=(",", ":")).encode() + ).hexdigest(), + "total_drafted": request.py_total_draft_tokens, + "total_accepted": request.py_total_accepted_draft_tokens, + } + + +@dataclass +class _ForwardJoin: + generation_id: int + executor_iteration: int + route: str + requests: tuple[Any, ...] + before: list[dict[str, Any]] + sample_index: int | None = None + update_before: list[dict[str, Any]] | None = None + sampler_update_applied: list[bool] | None = None + after: list[dict[str, Any]] | None = None + host_step_time_ms: float | None = None + + +class OwnedDSparkProfileEngine(DSparkCostProfileEngine): + """Only a source-owned executor boundary may enter a real request forward.""" + + def forward( + self, batch, resource_manager, new_tensors_device=None, cache_indirection_buffer=None + ): + if self.is_warmup: + return PyTorchModelEngine.forward( + self, + batch, + resource_manager, + new_tensors_device, + cache_indirection_buffer=cache_indirection_buffer, + ) + boundary = getattr(self, "_owned_forward_boundary", None) + if boundary is None or boundary[2] is not batch or boundary[3] is not new_tensors_device: + raise RuntimeError("No original-request owner boundary; ordinary serving denied") + route, cell, _, _ = boundary + if route == "cell": + return self.forward_profile_cell( + cell, + batch, + resource_manager, + new_tensors_device=new_tensors_device, + cache_indirection_buffer=cache_indirection_buffer, + ) + # Genuine prefill/first-draft/short tail/peer-not-ready transitions use + # the unchanged native path and produce no profile event record. + for request in batch.generation_requests: + request.py_verify_len = None + self.cuda_graph_runner.agreed_ragged_bucket = None + return PyTorchModelEngine.forward( + self, + batch, + resource_manager, + new_tensors_device, + cache_indirection_buffer=cache_indirection_buffer, + ) + + +class OwnedDSparkProfileExecutor(PyExecutor): + """Original executor loops plus fixed-cell provenance; not a serving factory.""" + + def __init__(self, *args, **kwargs): + self._owned_generation = None + self._owned_generation_ids = set() + self._owned_joins = [] + self._owned_pending = {} + self._owned_last_join = None + self._owned_request_objects = {} + self._owned_final_requests = {} + self._owned_failures = [] + self._owned_kv_capacity_enqueue = None + super().__init__(*args, **kwargs) + + @contextmanager + def _profiler(self): + """Join the original complete-loop timer, without new events or waits. + + The parent step reads the just-finished loop at the next loop boundary. + An overlap sampler may consume another batch there; its forward-event + duration and iterLatencyMS are not substitutes for this loop's cost. + """ + with super()._profiler() as step: + if self.is_warmup or self.model_engine._cost_profile.measurement_scope != "whole-step": + yield step + return + if not self.enable_iter_perf_stats: + raise ValueError("Whole-step costs require existing enable_iter_perf_stats") + previous_iteration = None + + def advance(): + nonlocal previous_iteration + step() + join = self._owned_joins[-1] if self._owned_joins else None + if ( + join is not None + and join.executor_iteration == previous_iteration + and join.route == "cell" + ): + duration = self._latest_host_step_time_ms + if ( + isinstance(duration, bool) + or not isinstance(duration, (int, float)) + or not math.isfinite(duration) + or duration <= 0 + ): + raise ValueError("Original cell loop has no finite positive hostStepTimeMS") + if join.host_step_time_ms is not None: + raise RuntimeError("Original cell loop timing was joined twice") + join.host_step_time_ms = duration + previous_iteration = self.iter_counter + + yield advance + + def begin_profile_generation(self, generation_id: int, cell: Any) -> None: + if type(generation_id) is not int or generation_id < 0: + raise ValueError("Original protocol generation ID must be a nonnegative integer") + self.model_engine._cost_profile.require_cell(cell) + with self.control_action(drain=True): + if self._owned_pending or generation_id in self._owned_generation_ids: + raise RuntimeError("Undrained or repeated original generation") + if self._owned_generation is not None: + raise RuntimeError("End the previous original generation before another") + self._owned_generation_ids.add(generation_id) + self._owned_generation = (generation_id, cell) + + def enqueue_requests(self, requests, result_wait_queue=None): + if self._owned_generation is None and not self.is_warmup: + authorization = self._owned_kv_capacity_enqueue + if ( + authorization is None + or requests is not authorization[0] + or authorization[1] + or result_wait_queue is not None + ): + raise RuntimeError("Ordinary requests require an explicit offline generation") + # The original KV estimator enqueues this exact internal list before + # setting warmup. Consume before delegation, even if enqueue fails. + self._owned_kv_capacity_enqueue = (requests, True) + return super().enqueue_requests(requests, result_wait_queue) + + def enqueue_request( + self, request: ExecutorRequest, result_wait_queue: "Optional[ActorHandle]" = None + ) -> int: + if self._owned_generation is None and not self.is_warmup: + raise RuntimeError("Ordinary requests require an explicit offline generation") + return super().enqueue_request(request, result_wait_queue) + + def _handle_errors( + self, + error_msg=None, + *, + requests=None, + charge_budget=True, + fatal_is_collective_aligned=False, + ): + if self._owned_generation is not None: + self._owned_failures.append(str(error_msg or "Original executor error")) + return super()._handle_errors( + error_msg, + requests=requests, + charge_budget=charge_budget, + fatal_is_collective_aligned=fatal_is_collective_aligned, + ) + + def _forward_step(self, batch, new_tensors_device=None): + if self.is_warmup: + return super()._forward_step(batch, new_tensors_device) + generation = self._owned_generation + engine = self.model_engine + requests = tuple(request for request in batch.all_requests() if not request.is_dummy) + ready = generation is not None + generation_id, cell = generation if ready else (-1, None) + eligible = ( + ready + and batch.num_context_requests == 0 + and batch.num_encoder_requests == 0 + and len(batch.generation_requests) == cell.graph_batch_size + and all( + not r.is_dummy + and not r.py_is_first_draft + and r.py_beam_width == 1 + and r.py_seq_slot is not None + for r in batch.generation_requests + ) + and len({r.py_request_id for r in batch.generation_requests}) + == len(batch.generation_requests) + and _valid_previous_state( + new_tensors_device, batch.generation_requests, engine._cost_profile.physical_k + ) + ) + row = [ + int(ready), + generation_id, + self.iter_counter, + cell.graph_batch_size if ready else -1, + cell.verifier_budget if ready else -1, + int(eligible), + ] + rows = self.dist.tp_allgather_int64(row).tolist() if self.dist.tp_size > 1 else [row] + if ( + len(rows) != self.dist.tp_size + or any(len(r) != 6 or r[0] != 1 for r in rows) + or any(r[:5] != rows[0][:5] for r in rows) + ): + raise RuntimeError("All ranks must agree on the original generation/iteration/cell") + route = "cell" if all(r[5] == 1 for r in rows) else "native-transition" + for request in requests: + key = (generation_id, request.py_request_id) + previous = self._owned_request_objects.setdefault(key, request) + if previous is not request: + raise RuntimeError("Request ID reused with a different original incarnation") + join = _ForwardJoin( + generation_id, + self.iter_counter, + route, + requests, + [_token_snapshot(request) for request in requests], + ) + before_records = len(engine.cuda_graph_runner._profile_records) + if getattr(engine, "_owned_forward_boundary", None) is not None: + raise RuntimeError("Overlapping owner forward boundaries") + engine._owned_forward_boundary = (route, cell, batch, new_tensors_device) + try: + output = super()._forward_step(batch, new_tensors_device) + if output is None: + raise RuntimeError("Original forward failed; no observation promotion") + count = len(engine.cuda_graph_runner._profile_records) + if count != before_records + int(route == "cell"): + raise RuntimeError("Native transition or fixed-cell replay count differs") + if route == "cell": + record = engine.cuda_graph_runner._profile_records[-1] + if tuple(record["request_ids"]) != tuple( + r.py_request_id for r in batch.generation_requests + ): + raise RuntimeError("Captured replay and original request identities differ") + join.sample_index = record["sample_index"] + self._owned_joins.append(join) + self._owned_last_join = join + return output + except BaseException as error: + self._owned_failures.append(type(error).__name__ + ": " + str(error)) + raise + finally: + engine._owned_forward_boundary = None + + def _sample_async(self, batch, output): + state = super()._sample_async(batch, output) + if self.is_warmup: + return state + join = self._owned_last_join + if state is None or join is None or id(state) in self._owned_pending: + raise RuntimeError("Missing or reused original sampler state") + if any( + not any(request is original for original in join.requests) + for request in state.requests + if not request.is_dummy + ): + raise RuntimeError("Sampler state belongs to another original forward") + if join.sample_index is not None and tuple(state.requests) != join.requests: + raise RuntimeError("Measured sampler requests differ from original decode forward") + self._owned_pending[id(state)] = (state, join) + self._owned_last_join = None + return state + + def _update_requests(self, state, resource_manager=None): + if self.is_warmup: + return super()._update_requests(state, resource_manager) + pending = self._owned_pending.get(id(state)) + if pending is None or pending[0] is not state: + raise RuntimeError("Unknown or already consumed original sampler state") + join = pending[1] + join.update_before = [_token_snapshot(request) for request in join.requests] + completed_before_update = [ + request.state == LlmRequestState.GENERATION_COMPLETE for request in join.requests + ] + super()._update_requests(state, resource_manager) + if join.sample_index is not None: + self.model_engine.cuda_graph_runner.complete_sample(join.sample_index) + # Read only after the original sampler sync/token/KV/stop update and + # original accumulation. Under overlap this is the previous iteration. + join.after = [_token_snapshot(request) for request in join.requests] + join.sampler_update_applied = [ + after["decoding_iter"] > before["decoding_iter"] + for before, after in zip(join.update_before, join.after) + ] + if join.sample_index is not None and any( + not applied and not completed + for applied, completed in zip(join.sampler_update_applied, completed_before_update) + ): + self._owned_failures.append("Original measured request update did not advance") + raise RuntimeError("Original measured request update did not advance") + del self._owned_pending[id(state)] + + def _free_request_resources(self, request): + generation = self._owned_generation + final = None + if generation is not None and not request.is_dummy: + key = (generation[0], request.py_request_id) + if ( + self._owned_request_objects.get(key) is request + and request.state == LlmRequestState.GENERATION_COMPLETE + ): + final = (key, _token_snapshot(request)) + result = super()._free_request_resources(request) + if final is not None: + self._owned_final_requests[final[0]] = final[1] + return result + + def _profile_drain_predicates(self): + previous = self.previous_batch + state = previous.sample_state if previous is not None else None + pending = self._owned_pending.get(id(state)) if state is not None else None + join = pending[1] if pending is not None else None + return { + "generation_id": self._owned_generation[0] if self._owned_generation else None, + "failures": len(self._owned_failures), + "active_requests": len(self.active_requests), + "waiting_requests": len(self.waiting_queue), + "pending_samples": len(self._owned_pending), + "last_join_present": self._owned_last_join is not None, + "forward_boundary_present": ( + getattr(self.model_engine, "_owned_forward_boundary", None) is not None + ), + "previous_batch_present": previous is not None, + "pending_matches_previous": ( + state is not None + and len(self._owned_pending) == 1 + and pending is not None + and pending[0] is state + and any(join is original for original in self._owned_joins) + and self._owned_generation is not None + and join.generation_id == self._owned_generation[0] + and join.update_before is None + and join.sampler_update_applied is None + and join.after is None + ), + } + + def end_profile_generation(self, *, control=None) -> dict[str, Any]: + if self._owned_generation is None: + raise RuntimeError("No original generation to drain") + with self.control_action(drain=True): + before = self._profile_drain_predicates() + if control is None: + raise RuntimeError( + "Offline profile drain needs the original gang-control communicator" + ) + peers = control.allgather(before) + quiescent = all( + row["failures"] == 0 + and row["active_requests"] == 0 + and row["waiting_requests"] == 0 + and not row["last_join_present"] + and not row["forward_boundary_present"] + and row["generation_id"] == before["generation_id"] + for row in peers + ) + no_tail = all(row["pending_samples"] == 0 for row in peers) + exact_tail = all( + row["pending_samples"] == 1 and row["pending_matches_previous"] for row in peers + ) + if not quiescent or not (no_tail or exact_tail): + raise RuntimeError( + f"Original generation has not cleanly drained: gang_predicates={peers!r}" + ) + consumed = False + if exact_tail: + # All ranks are paused in the original exclusive control action. + # Delegate every sampler/update/resource/timing operation to the + # existing executor helper; do not clear pending joins directly. + self._consume_previous_batch_for_rebalance() + consumed = True + after = self._profile_drain_predicates() + if ( + self._owned_pending + or self._owned_last_join is not None + or self._owned_failures + or self.active_requests + or len(self.waiting_queue) + ): + raise RuntimeError( + "Original generation has not cleanly drained: " + f"before={before!r}, after={after!r}" + ) + generation_id = self._owned_generation[0] + keys = [key for key in self._owned_request_objects if key[0] == generation_id] + if not keys or any(key not in self._owned_final_requests for key in keys): + raise RuntimeError("Original request final resource-release joins are incomplete") + observation = self.model_engine.profile_observations() + by_sample = { + join.sample_index: join + for join in self._owned_joins + if join.sample_index is not None + } + if set(by_sample) != {row["sample_index"] for row in observation["samples"]}: + raise RuntimeError("Completed replay events and lifecycle joins differ") + joins = [ + { + "generation_id": join.generation_id, + "executor_iteration": join.executor_iteration, + "route": join.route, + "sample_index": join.sample_index, + "before": join.before, + "update_before": join.update_before, + "sampler_update_applied": join.sampler_update_applied, + "after": join.after, + "host_step_time_ms": join.host_step_time_ms, + } + for join in self._owned_joins + if join.generation_id == generation_id + ] + if any(join["after"] is None for join in joins): + raise RuntimeError("Original sampling/update join missing") + if self.model_engine._cost_profile.measurement_scope == "whole-step" and any( + join["sample_index"] is not None and join["host_step_time_ms"] is None + for join in joins + ): + raise RuntimeError("Measured cell lacks its original complete-loop timing") + generation_samples = { + join["sample_index"] for join in joins if join["sample_index"] is not None + } + if not generation_samples: + raise RuntimeError("Prescribed generation did not execute a measured graph cell") + observation = { + **observation, + "samples": [ + row + for row in observation["samples"] + if row["sample_index"] in generation_samples + ], + } + self._owned_generation = None + return { + "schema": "dspark-original-request-profile-generation-v1", + "generation_id": generation_id, + "drain": { + "before": before, + "after": after, + "original_previous_batch_consumed": consumed, + }, + "observation": observation, + "lifecycle_joins": joins, + "original_final_requests": [self._owned_final_requests[key] for key in keys], + "original_shutdown_still_required": True, + "task_accuracy_qualified": False, + "native_or_gpu_qualified": False, + "full_policy_h": None, + "production_admission_qualified": False, + } + + +class _ProfileKvCacheCreator(KvCacheCreator): + def configure_kv_cache_capacity(self, py_executor=None) -> None: + if py_executor is None or self._skip_est: + return super().configure_kv_cache_capacity(py_executor) + if ( + not isinstance(py_executor, OwnedDSparkProfileExecutor) + or py_executor._owned_generation is not None + or py_executor._owned_kv_capacity_enqueue is not None + ): + raise RuntimeError("Invalid or nested offline KV startup phase") + py_executor._owned_kv_capacity_enqueue = (self._dummy_reqs, False) + try: + return super().configure_kv_cache_capacity(py_executor) + finally: + py_executor._owned_kv_capacity_enqueue = None + + +def create_profile_executor(*, llm_args, profile, checkpoint_dir=None, tokenizer=None): + """Construct only the prescribed offline owner; never a serving executor. + + No source pins, private manifests, historical receipts or global cloning are + required. Runtime and loaded-byte identity is measured by the public driver. + Normal factory construction, resource/KV estimation and loops are retained. + """ + if llm_args.enable_resource_governor: + raise ValueError("The offline profile owner does not support resource-governor control") + constructors = _RuntimeConstructors( + model_engine=partial(OwnedDSparkProfileEngine, profile=profile), + executor=OwnedDSparkProfileExecutor, + kv_cache_creator=_ProfileKvCacheCreator, + ) + return create_py_executor( + llm_args=llm_args, + checkpoint_dir=checkpoint_dir, + tokenizer=tokenizer, + runtime_constructors=constructors, + ) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 2d25b76c30d1..6931a0c36df4 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -629,13 +629,17 @@ def _init_decoder_state(self) -> None: self._dynamic_draft_len_mapping = self._compute_dynamic_draft_len_mapping() def _initialize_cuda_graph_runner(self) -> CUDAGraphRunner | None: - runner_type = DSparkCUDAGraphRunner if self._dspark_confidence_enabled else CUDAGraphRunner - cuda_graph_runner = runner_type(self._cuda_graph_runner_config()) + cuda_graph_runner = self._create_cuda_graph_runner(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) return cuda_graph_runner + def _create_cuda_graph_runner(self, config: CUDAGraphRunnerConfig) -> CUDAGraphRunner: + """Construction seam for explicit offline tooling, with ordinary defaults retained.""" + runner_type = DSparkCUDAGraphRunner if self._dspark_confidence_enabled else CUDAGraphRunner + return runner_type(config) + def _cuda_graph_runner_config(self) -> CUDAGraphRunnerConfig: return CUDAGraphRunnerConfig( use_cuda_graph=( @@ -3783,12 +3787,12 @@ def _apply_device_window_prologue(self, inputs, new_tensors_device) -> bool: 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_schedule import DSparkFusedScheduleError from tensorrt_llm._torch.speculative.dspark_select import ( gather_packed_draft_tokens, materialize_compact_layout, select_windows_device, ) - from tensorrt_llm._torch.speculative.dspark_schedule import DSparkFusedScheduleError runner = self.cuda_graph_runner budget = self._dspark_device_budget diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 348328d75d96..83a3e3cbd00b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -11,7 +11,7 @@ from contextlib import contextmanager from dataclasses import dataclass from itertools import chain -from typing import Optional, Tuple +from typing import Callable, Optional, Tuple import torch from strenum import StrEnum @@ -69,6 +69,19 @@ FP4_MLA_TOKENS_PER_BLOCK = 128 +@dataclass(frozen=True) +class _RuntimeConstructors: + """Explicit internal constructors for a non-serving offline executor. + + Normal serving passes no bundle and retains the ordinary constructors. + This is not an LLM argument and cannot enable confidence admission. + """ + + model_engine: Callable[..., PyTorchModelEngine] + executor: type[PyExecutor] + kv_cache_creator: type[KvCacheCreator] + + class _ExecutorMemoryMonitor: """Currently this focuses on tracking memory usage and related errors.""" @@ -484,6 +497,7 @@ def _create_py_executor_impl( tokenizer: Optional[TokenizerBase] = None, profiling_stage_data: Optional[dict] = None, resource_governor_queue=None, + runtime_constructors: Optional[_RuntimeConstructors] = None, ) -> PyExecutor: """Create and initialize a PyExecutor instance from the given LLM arguments. @@ -515,6 +529,7 @@ def _create_py_executor_impl( tokenizer=tokenizer, profiling_stage_data=profiling_stage_data, resource_governor_queue=resource_governor_queue, + runtime_constructors=runtime_constructors, ) py_executor.metrics.update(creation_metrics) return py_executor @@ -528,6 +543,7 @@ def _create_py_executor( tokenizer: Optional[TokenizerBase] = None, profiling_stage_data: Optional[dict] = None, resource_governor_queue=None, + runtime_constructors: Optional[_RuntimeConstructors] = None, ) -> PyExecutor: """Create and initialize a PyExecutor while recording scoped metrics.""" initial_model_engine_metrics: dict[str, float | dict[str, float]] = {} @@ -790,7 +806,9 @@ def allocation_scope(current_stage: ExecutorMemoryType): model_weights_restore_mode = sleep_config.restore_modes[ ExecutorMemoryType.MODEL_WEIGHTS_MAIN] - model_engine = PyTorchModelEngine( + model_engine_type = (PyTorchModelEngine if runtime_constructors is None + else runtime_constructors.model_engine) + model_engine = model_engine_type( model_path=checkpoint_dir, llm_args=llm_args, mapping=mapping, @@ -1130,7 +1148,9 @@ def allocation_scope(current_stage: ExecutorMemoryType): # Get draft config for one-engine speculative decoding if available draft_config = getattr(model_engine.model, 'draft_config', None) - kv_cache_creator = KvCacheCreator( + kv_cache_creator_type = (KvCacheCreator if runtime_constructors is None + else runtime_constructors.kv_cache_creator) + kv_cache_creator = kv_cache_creator_type( model_engine=model_engine, draft_model_engine=draft_model_engine, mapping=mapping, @@ -1196,6 +1216,8 @@ def allocation_scope(current_stage: ExecutorMemoryType): # run gc.collect() to free memory of the previous py_executor, avoid cudaFree overlap with cuda graph capture gc.collect() py_executor = create_py_executor_instance( + executor_type=None + if runtime_constructors is None else runtime_constructors.executor, dist=dist, resources=resources, mapping=mapping, @@ -1277,6 +1299,8 @@ def allocation_scope(current_stage: ExecutorMemoryType): # run gc.collect() to free memory of the previous py_executor, avoid cudaFree overlap with cuda graph capture gc.collect() py_executor = create_py_executor_instance( + executor_type=None if runtime_constructors is None else + runtime_constructors.executor, dist=dist, resources=resources, mapping=mapping, @@ -1341,6 +1365,7 @@ def create_py_executor( tokenizer: Optional[TokenizerBase] = None, profiling_stage_data: Optional[dict] = None, resource_governor_queue=None, + runtime_constructors: Optional[_RuntimeConstructors] = None, ) -> PyExecutor: """Create a PyExecutor and roll back a partially initialized DWDP runtime.""" previous_dwdp_manager = get_global_dwdp_manager() @@ -1354,6 +1379,7 @@ def create_py_executor( tokenizer=tokenizer, profiling_stage_data=profiling_stage_data, resource_governor_queue=resource_governor_queue, + runtime_constructors=runtime_constructors, ) except BaseException: current_dwdp_manager = get_global_dwdp_manager() diff --git a/tensorrt_llm/_torch/speculative/dspark_cost_profile.py b/tensorrt_llm/_torch/speculative/dspark_cost_profile.py new file mode 100644 index 000000000000..db8a9ded3e96 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_cost_profile.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Explicit non-serving DSpark graph-cell profiling contracts.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Literal + +from pydantic import ConfigDict, Field, StrictBool, StrictInt, model_validator + +from tensorrt_llm.llmapi.utils import StrictBaseModel + +__all__ = ["DSparkCostProfileConfig", "DSparkProfileCell"] + + +class DSparkProfileCell(StrictBaseModel): + """One rank-local graph shape; zero budget denotes the native full block.""" + + model_config = ConfigDict(frozen=True) + graph_batch_size: StrictInt = Field(gt=0, description="Rank-local request count G.") + verifier_budget: StrictInt = Field( + ge=0, description="Verifier token count V; zero means native." + ) + + +class DSparkCostProfileConfig(StrictBaseModel): + """Construction-only graph grid, never a serving admission configuration.""" + + model_config = ConfigDict(frozen=True) + physical_k: StrictInt = Field(default=3, gt=0, description="Full physical draft block K.") + cells: tuple[DSparkProfileCell, ...] = Field(min_length=1, description="Declared graph cells.") + record_confidence_head: StrictBool = Field( + default=True, description="Execute the head in graphs." + ) + measurement_scope: Literal["graph-replay", "whole-step"] = Field( + default="graph-replay", description="Original replay or executor-loop cost boundary." + ) + max_samples: StrictInt = Field( + default=4096, gt=0, le=65536, description="Maximum raw replay records for the run." + ) + + @model_validator(mode="after") + def _validate_grid(self) -> DSparkCostProfileConfig: + seen = set() + native = set() + for cell in self.cells: + pair = (cell.graph_batch_size, cell.verifier_budget) + if pair in seen: + raise ValueError("Profile cells must not contain duplicates") + seen.add(pair) + if cell.verifier_budget == 0: + native.add(cell.graph_batch_size) + elif ( + self.physical_k == 1 + or not cell.graph_batch_size + <= cell.verifier_budget + < cell.graph_batch_size * (self.physical_k + 1) + ): + raise ValueError("Compact cells require K>1 and G <= V < G*(K+1); native uses V=0") + if native != {cell.graph_batch_size for cell in self.cells}: + raise ValueError("Every profiled G requires a V=0 native control") + return self + + def validate_graph_ladder(self, graph_batch_sizes: Sequence[int]) -> None: + """Require the exact runtime ladder, not an undocumented subset.""" + if any(type(value) is not int or value <= 0 for value in graph_batch_sizes): + raise ValueError("Runtime graph row counts must be positive integers") + if len(set(graph_batch_sizes)) != len(graph_batch_sizes): + raise ValueError("Runtime graph ladder contains duplicates") + if set(graph_batch_sizes) != {cell.graph_batch_size for cell in self.cells}: + raise ValueError("Profile grid must cover the exact captured graph ladder") + + def graph_shapes(self, graph_batch_sizes: Sequence[int]) -> list[tuple[int, ...]]: + """Return native and positive-V capture shapes in deterministic order.""" + self.validate_graph_ladder(graph_batch_sizes) + return [ + (cell.graph_batch_size, self.physical_k) + if cell.verifier_budget == 0 + else (cell.graph_batch_size, self.physical_k, cell.verifier_budget) + for cell in sorted( + self.cells, key=lambda item: (item.graph_batch_size, item.verifier_budget) + ) + ] + + def require_cell(self, cell: DSparkProfileCell) -> None: + """Reject an undeclared shape before request metadata can change.""" + if cell not in self.cells: + raise ValueError("Requested profile cell is not in the construction grid") diff --git a/tensorrt_llm/_torch/speculative/dspark_cost_profiler.py b/tensorrt_llm/_torch/speculative/dspark_cost_profiler.py new file mode 100644 index 000000000000..1b3da52abfc4 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_cost_profiler.py @@ -0,0 +1,563 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSpark chain adapter and explicit offline driver; no serving admission.""" + +from __future__ import annotations + +import hashlib +import importlib +import importlib.metadata +import math +import sys +import time +from functools import partial +from pathlib import Path +from typing import Callable + +import torch +import yaml +from mpi4py import MPI + +from tensorrt_llm._torch.pyexecutor.dspark_request_profile import create_profile_executor +from tensorrt_llm._torch.speculative.dspark_cost_profile import DSparkCostProfileConfig +from tensorrt_llm._utils import mpi_comm +from tensorrt_llm.executor.request import GenerationRequest +from tensorrt_llm.executor.worker import GenerationExecutorWorker +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig, TorchLlmArgs +from tensorrt_llm.sampling_params import SamplingParams +from tensorrt_llm.tools.dspark_cost_export import ( + observed_engine_fingerprint, + source_origin_from_args, + source_origin_input_metadata, +) +from tensorrt_llm.tools.verify_costs import ( + canonical_sha256, + create_private_output_directory, + isolated_control_communicator, + parse_json, + sha256_file, + write_new_json, +) + +__all__ = ["run_profile"] + + +class OwnedInputWorker(GenerationExecutorWorker): + """Construct the owned factory inside each original worker, not its proxy.""" + + def __init__( + self, + *, + profile: DSparkCostProfileConfig, + **kwargs, + ) -> None: + self._profile = profile + if kwargs.get("rpc_addr") is not None: + raise ValueError("Offline profiling cannot expose a serving RPC endpoint") + super().__init__(**kwargs) + + def _get_pytorch_executor_factory(self) -> Callable: + # Retain every original worker setup operation and its ordering. + return partial(create_profile_executor, profile=self._profile) + + +def response_rows(worker, inputs, output_length, timeout): + """Use original request conversion/enqueue/futures, with real native IDs.""" + deadline = time.monotonic() + timeout + futures = [] + for row in inputs: + if len(row["prompt_token_ids"]) + output_length > worker.max_seq_len: + raise ValueError("Original max_seq_len would shorten the fixed output length") + sampling = SamplingParams( + max_tokens=output_length, + temperature=0.0, + top_k=1, + n=1, + best_of=1, + ignore_eos=True, + detokenize=False, + exclude_input_from_output=True, + ) + request = GenerationRequest( + prompt_token_ids=list(row["prompt_token_ids"]), + sampling_params=sampling, + streaming=False, + ) + if time.monotonic() >= deadline: + raise TimeoutError("Original total child request budget exhausted before enqueue") + future = worker.submit(request) + native_id = worker._client_id_to_request_id[request.id] + futures.append((row, request.id, native_id, future)) + results = [] + for row, client_id, native_id, future in futures: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("Original total child request budget exhausted before response") + result = future.result(timeout=remaining) + outputs = result.outputs + if ( + len(outputs) != 1 + or outputs[0].finish_reason != "length" + or len(outputs[0].token_ids) != output_length + or any(type(token) is not int or token < 0 for token in outputs[0].token_ids) + ): + raise RuntimeError("Original response did not complete the fixed greedy output") + results.append( + { + "input_id": row["input_id"], + "client_id": client_id, + "native_request_id": native_id, + "prompt_token_ids": row["prompt_token_ids"], + "output_token_ids": list(outputs[0].token_ids), + "finish_reason": outputs[0].finish_reason, + } + ) + return results + + +def _collective(control, label, callback): + error, result = None, None + try: + result = callback() + except Exception as exception: + error = type(exception).__name__ + ": " + str(exception) + errors = control.allgather(error) + if any(error is not None for error in errors): + raise RuntimeError(label + ": " + repr(errors)) + return result + + +def _identity(worker, args, profile, expected_files): + """Record live installed sources/runtime and adapter-defined execution layout.""" + modules = ( + "tensorrt_llm._torch.pyexecutor.model_engine", + "tensorrt_llm._torch.pyexecutor._util", + "tensorrt_llm._torch.pyexecutor.py_executor_creator", + "tensorrt_llm._torch.pyexecutor.cuda_graph_runner", + "tensorrt_llm._torch.pyexecutor.py_executor", + "tensorrt_llm.executor.base_worker", + "tensorrt_llm.executor.worker", + "tensorrt_llm.llmapi.llm_args", + "tensorrt_llm._torch.speculative.dspark", + "tensorrt_llm._torch.pyexecutor.dspark_profile_engine", + "tensorrt_llm._torch.pyexecutor.dspark_request_profile", + "tensorrt_llm._torch.speculative.dspark_cost_profile", + "tensorrt_llm._torch.speculative.dspark_cost_profiler", + "tensorrt_llm.tools.verify_costs", + ) + if profile.measurement_scope == "whole-step": + modules += ("tensorrt_llm.tools.dspark_cost_export",) + sources = {} + for name in modules: + module = importlib.import_module(name) + sources[name] = sha256_file(Path(module.__file__)) + native_paths = set(torch.ops.loaded_libraries) + for name, module in tuple(sys.modules.items()): + if name.startswith(("tensorrt_llm.", "torch._C")): + path = getattr(module, "__file__", None) + if path is not None and Path(path).suffix in (".so", ".pyd"): + native_paths.add(path) + native = {str(Path(path).resolve()): sha256_file(Path(path)) for path in sorted(native_paths)} + if not native: + raise RuntimeError("Loaded native library identity is missing") + package = importlib.metadata.metadata("tensorrt_llm") + source_urls = [ + url for url in package.get_all("Project-URL", []) if url.startswith("Source Commit,") + ] + device = torch.cuda.get_device_properties(torch.cuda.current_device()) + versions = { + "torch": str(torch.__version__), + "cuda": torch.version.cuda, + "tensorrt_llm": importlib.metadata.version("tensorrt_llm"), + } + actual_files = { + "inputs_sha256": sha256_file(args.inputs), + "model_config_sha256": sha256_file(args.model / "config.json"), + } + if profile.measurement_scope == "whole-step": + actual_files["recorded_source_origin"] = source_origin_from_args(args) + actual_files["source_origin_observation"] = source_origin_input_metadata( + getattr(args, "recorded_source_origin", None), + getattr(args, "recorded_source_origin_sha256", None), + ) + if actual_files != expected_files: + raise ValueError("Profile inputs or model configuration changed after validation") + mapping = worker.engine.model_engine.mapping + model = worker.engine.model_engine.model + attention_classes = sorted( + { + type(module).__module__ + "." + type(module).__qualname__ + for module in model.modules() + if type(module).__module__.startswith("tensorrt_llm._torch.attention") + } + ) + if profile.measurement_scope == "whole-step": + # Retain observed model/module overlays as well as the adapter roster. + # Native library bytes are already recorded independently above. + for instance in model.modules(): + name = type(instance).__module__ + if name.startswith("tensorrt_llm.") and name not in sources: + module = importlib.import_module(name) + path = getattr(module, "__file__", None) + if path is not None and Path(path).suffix == ".py": + sources[name] = sha256_file(Path(path)) + return { + "identity_scope": "DIAGNOSTIC_PARTIAL_LOADED_RUNTIME", + "production_live_fingerprint_qualified": False, + "loaded_source_sha256": sources, + "package_source_commit_urls": source_urls, + "runtime": versions, + "local_native_libraries": native, + "local_device": { + "name": device.name, + "major": device.major, + "minor": device.minor, + "total_memory": device.total_memory, + }, + "topology": { + "tp": mapping.tp_size, + "pp": mapping.pp_size, + "cp": mapping.cp_size, + "ep": mapping.moe_ep_size, + "attention_dp": mapping.enable_attention_dp, + }, + "runtime_configuration": { + "max_batch_size": worker.llm_args.max_batch_size, + "max_seq_len": worker.max_seq_len, + "max_num_tokens": worker.llm_args.max_num_tokens, + "disable_overlap_scheduler": worker.llm_args.disable_overlap_scheduler, + "enable_iter_perf_stats": worker.llm_args.enable_iter_perf_stats, + "iter_perf_stats_interval": worker.llm_args.iter_perf_stats_interval, + "cuda_graph_config": worker.llm_args.cuda_graph_config.model_dump(mode="json"), + "kv_cache_config": worker.llm_args.kv_cache_config.model_dump(mode="json"), + "configured_attention_backend": str(worker.llm_args.attn_backend), + }, + **actual_files, + "profile": profile.model_dump(mode="json"), + "generations_per_shape": args.generations, + "sampling": { + "output_tokens": args.output_tokens, + "temperature": 0.0, + "top_k": 1, + "ignore_eos": True, + "beam_width": 1, + }, + "declared_shape_ids": [canonical_sha256(_shape(cell, profile)) for cell in profile.cells], + "observed_attention_module_classes": attention_classes, + "kernel_dispatch_qualification": False, + } + + +def _shape(cell, profile) -> dict: + """DSpark-specific shape identity; never reduce different attention layouts.""" + return { + "workload": "dspark_chain", + "physical_k": profile.physical_k, + "graph_batch_size": cell.graph_batch_size, + "verifier_budget": cell.verifier_budget, + "attention_layout": "uniform" if cell.verifier_budget == 0 else "packed_chain_rows", + "captured_head": profile.record_confidence_head, + } + + +def _raw_generation(generation, cell, repetition, identity, rank, world): + """The DSpark adapter decodes lifecycle records; the shared reducer does not.""" + samples = {row["sample_index"]: row for row in generation["observation"]["samples"]} + output = [] + for join in generation["lifecycle_joins"]: + index = join["sample_index"] + if index is None: + continue + sample = samples.pop(index) + scope = identity["profile"].get("measurement_scope", "graph-replay") + duration = ( + join["host_step_time_ms"] + if scope == "whole-step" + else sample["replay_entry_elapsed_ms"] + ) + output.append( + { + "sample_index": index, + "executor_iteration": join["executor_iteration"], + "elapsed_ms": duration, + "graph_replay_diagnostic_ms": sample["replay_entry_elapsed_ms"], + "eligible": all(join["sampler_update_applied"]), + "request_ids": sample["request_ids"], + "graph_key": sample["graph_key"], + } + ) + if samples: + raise RuntimeError("Unjoined graph replay samples") + graph_keys = {canonical_sha256(row["graph_key"]) for row in output} + if len(graph_keys) != 1: + raise RuntimeError("Prescribed cell observed zero or incompatible graph replay keys") + profile = DSparkCostProfileConfig(**identity["profile"]) + shape = _shape(cell, profile) + return { + "schema": "verify-cost-raw-v1", + "rank": rank, + "world_size": world, + "generation": repetition, + "generation_closed": True, + "model_shutdown_completed": False, + "identity": identity, + "timing_scope": ( + "original-executor-loop-hostStepTimeMS" + if profile.measurement_scope == "whole-step" + else "original-graph-replay-entry-including-static-copies" + ), + "declared_shape_id": canonical_sha256(shape), + "execution_shape": {**shape, "observed_graph_key": output[0]["graph_key"]}, + "samples": output, + "lifecycle": generation, + "policy_overhead_measured": False, + "production_admission_qualified": False, + } + + +def _prepare(args, world): + """Validate rank-local inputs before any peer enters collective construction.""" + if world > 1 and MPI.Query_thread() < MPI.THREAD_MULTIPLE: + raise ValueError("Concurrent runtime/profile MPI requires THREAD_MULTIPLE") + if ( + args.generations < 1 + or args.output_tokens < 1 + or not math.isfinite(args.timeout) + or args.timeout <= 0 + ): + raise ValueError("Positive generations, output length and total timeout required") + grid = parse_json(args.grid.read_bytes()) + if not isinstance(grid, dict): + raise ValueError("The profile grid must be a JSON object") + if "physical_k" in grid and grid["physical_k"] != args.physical_k: + raise ValueError("CLI physical K and grid differ") + cost_scope = getattr(args, "cost_scope", "graph-replay") + if "measurement_scope" in grid and grid["measurement_scope"] != cost_scope: + raise ValueError("CLI measurement scope and grid differ") + profile = DSparkCostProfileConfig( + **{**grid, "physical_k": args.physical_k, "measurement_scope": cost_scope} + ) + config = yaml.safe_load(args.config.read_text(encoding="utf-8")) + if not isinstance(config, dict): + raise ValueError("The runtime configuration must be a YAML mapping") + llm_args = TorchLlmArgs(**{**config, "model": str(args.model), "backend": "pytorch"}) + if profile.measurement_scope == "whole-step" and ( + not llm_args.enable_iter_perf_stats or getattr(args, "source_root", None) is None + ): + raise ValueError("Whole-step collection requires enable_iter_perf_stats and --source-root") + spec = llm_args.speculative_config + mapping = llm_args.parallel_config.to_mapping() + if profile.measurement_scope == "whole-step" and world > 1 and not mapping.enable_attention_dp: + raise ValueError( + "Schema-v2 whole-step export currently requires attention-DP fanout or one rank" + ) + if ( + not isinstance(spec, DSparkDecodingConfig) + or spec.max_draft_len != profile.physical_k + or mapping.tp_size != world + or mapping.pp_size != 1 + or mapping.cp_size != 1 + ): + raise ValueError("Actual MPI topology/model physical K must match the profile") + if ( + llm_args.guided_decoding_backend is not None + or llm_args.batched_logits_processor is not None + or llm_args.lora_config is not None + or llm_args.cache_transceiver_config is not None + or llm_args.kv_connector_config is not None + or llm_args.enable_in_graph_sampling + or llm_args.encode_only + or llm_args.enable_resource_governor + ): + raise ValueError("The offline driver requires aggregate greedy generation without adapters") + input_bytes = args.inputs.read_bytes() + inputs = [parse_json(line) for line in input_bytes.splitlines() if line.strip()] + factor = world if mapping.enable_attention_dp else 1 + maximum = max(cell.graph_batch_size for cell in profile.cells) * factor + if len(inputs) < maximum or any( + not isinstance(row, dict) + or not isinstance(row.get("prompt_token_ids"), list) + or not row["prompt_token_ids"] + or any(type(token) is not int or token < 0 for token in row["prompt_token_ids"]) + for row in inputs + ): + raise ValueError("Enough nonempty original integer-token prompts are required") + for index, row in enumerate(inputs): + row["input_id"] = index + files = { + "inputs_sha256": hashlib.sha256(input_bytes).hexdigest(), + "model_config_sha256": sha256_file(args.model / "config.json"), + } + if profile.measurement_scope == "whole-step": + files["recorded_source_origin"] = source_origin_from_args(args) + files["source_origin_observation"] = source_origin_input_metadata( + getattr(args, "recorded_source_origin", None), + getattr(args, "recorded_source_origin_sha256", None), + ) + fingerprint = canonical_sha256( + { + "config": config, + "profile": profile.model_dump(mode="json"), + "inputs": inputs, + "generations": args.generations, + "output_tokens": args.output_tokens, + "timeout": args.timeout, + "output_directory": str(args.output.resolve()), + **files, + } + ) + return profile, llm_args, inputs, factor, fingerprint, files + + +def run_profile(args) -> None: + """Use original worker requests and state updates for a prescribed offline grid.""" + # Original request-count/payload broadcasts stay on mpi_comm(). Main-thread + # profile collectives need a separate context before any worker can start. + with isolated_control_communicator(mpi_comm()) as control: + _run_profile(args, control) + + +def _run_profile(args, control) -> None: + """Keep the isolated control context alive through startup and all closure.""" + rank, world = control.Get_rank(), control.Get_size() + profile, llm_args, inputs, factor, fingerprint, expected_files = _collective( + control, "prepare", lambda: _prepare(args, world) + ) + fingerprints = control.allgather(fingerprint) + if len(set(fingerprints)) != 1: + raise ValueError("Ranks disagree on inputs/configuration/grid") + _collective( + control, + "output directory", + lambda: create_private_output_directory(args.output) if rank == 0 else None, + ) + worker = None + records = [] + succeeded = False + shutdown_completed = False + try: + + def construct(): + nonlocal worker + # Preserve a successful local owner even if a peer reports failure + # at the following collective error-agreement boundary. + worker = OwnedInputWorker( + profile=profile, + engine=None, + hf_model_dir=args.model, + tokenizer=None, + llm_args=llm_args, + is_llm_executor=True, + ) + return worker + + _collective(control, "construct", construct) + identity = _collective( + control, "identity", lambda: _identity(worker, args, profile, expected_files) + ) + identity["devices"] = control.allgather(identity.pop("local_device")) + identity["rank_native_libraries"] = control.allgather( + identity.pop("local_native_libraries") + ) + if profile.measurement_scope == "whole-step": + identity["recorded_engine_fingerprint"] = _collective( + control, + "recorded cost engine metadata", + lambda: observed_engine_fingerprint( + identity, + source_root=args.source_root, + world_size=world, + recorded_origin=getattr(args, "recorded_source_origin", None), + recorded_origin_sha256=getattr(args, "recorded_source_origin_sha256", None), + ), + ) + if len(set(control.allgather(canonical_sha256(identity)))) != 1: + raise ValueError("Actual rank identities differ") + deadline = time.monotonic() + args.timeout + generation_id = 0 + for cell in profile.cells: + selected = inputs[: factor * cell.graph_batch_size] + for repetition in range(args.generations): + _collective( + control, + "begin", + lambda: worker.engine.begin_profile_generation(generation_id, cell), + ) + responses, error = None, None + if rank == 0: + try: + responses = response_rows( + worker, + selected, + args.output_tokens, + max(0, deadline - time.monotonic()), + ) + except Exception as exception: + error = type(exception).__name__ + ": " + str(exception) + error, responses = control.bcast((error, responses), root=0) + if error is not None: + raise RuntimeError(error) + generation = _collective( + control, + "drain", + lambda: worker.engine.end_profile_generation(control=control), + ) + + def join_generation(): + by_id = {row["native_request_id"]: row for row in responses} + if len(by_id) != len(responses): + raise RuntimeError("Duplicate native response IDs") + for final in generation["original_final_requests"]: + response = by_id[final["request_id"]] + if final["tokens"] != ( + response["prompt_token_ids"] + response["output_token_ids"] + ): + raise RuntimeError("Final native request tokens and response differ") + return _raw_generation(generation, cell, repetition, identity, rank, world) + + record = _collective(control, "generation join", join_generation) + record["responses"] = responses + records.append(record) + generation_id += 1 + succeeded = True + finally: + original_failure = sys.exc_info()[1] + try: + # Every rank enters the same control phase, including a rank whose + # constructor failed. Do not manufacture a model-shutdown receipt + # for a worker that was never successfully constructed. + _collective( + control, "shutdown", lambda: worker.shutdown() if worker is not None else None + ) + shutdown_completed = worker is not None + finally: + for index, record in enumerate(records): + record["model_shutdown_completed"] = shutdown_completed + record["run_complete"] = succeeded and shutdown_completed + _collective( + control, + "preserve raw", + lambda: write_new_json( + args.output / f"rank-{rank}-record-{index}.json", record + ), + ) + if not succeeded or not shutdown_completed: + _collective( + control, + "preserve failure", + lambda: write_new_json( + args.output / f"rank-{rank}-failure.json", + { + "schema": "verify-cost-run-failure-v1", + "run_complete": False, + "failure_type": ( + type(original_failure).__name__ + if original_failure + else "ShutdownFailure" + ), + "model_shutdown_completed": shutdown_completed, + }, + ), + ) + if not succeeded: + raise RuntimeError("Offline generation did not complete") diff --git a/tensorrt_llm/_torch/speculative/dspark_reproduction.py b/tensorrt_llm/_torch/speculative/dspark_reproduction.py new file mode 100644 index 000000000000..109339584460 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_reproduction.py @@ -0,0 +1,370 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSpark adapter for explicit offline collection and ordinary automatic policy.""" + +from __future__ import annotations + +import copy +import importlib +import importlib.metadata +import math +import os +import sys +import time +from pathlib import Path + +import torch +import yaml +from mpi4py import MPI + +from tensorrt_llm._torch.pyexecutor.dspark_observation_runtime import observation_factory +from tensorrt_llm._torch.speculative.dspark_cost_profiler import _collective, response_rows +from tensorrt_llm._utils import mpi_comm +from tensorrt_llm.executor.worker import GenerationExecutorWorker +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig, TorchLlmArgs +from tensorrt_llm.tools.dspark_cost_export import ( + source_origin_from_args, + source_origin_input_metadata, +) +from tensorrt_llm.tools.dspark_observations import validate_responses +from tensorrt_llm.tools.verify_costs import ( + canonical_sha256, + create_private_output_directory, + isolated_control_communicator, + parse_json, + sha256_file, + write_new_json, +) + +_POLICY_FIELDS = ( + "enable_confidence_scheduling", + "enable_fused_confidence_scheduler", + "validate_confidence_artifact_freshness", + "confidence_sts_path", + "confidence_sps_table_path", + "confidence_sps_live_fingerprint_path", + "confidence_admission_receipt_path", + "confidence_admission_receipt_sha256", +) + + +class ObservationWorker(GenerationExecutorWorker): + def __init__(self, *, collect_calibration=False, **kwargs): + self._collect_calibration = collect_calibration + if kwargs.get("rpc_addr") is not None: + raise ValueError("Offline collection/benchmark cannot expose serving RPC") + super().__init__(**kwargs) + + def _get_pytorch_executor_factory(self): + return observation_factory(collect_calibration=self._collect_calibration) + + +def _prepare(args, world): + if os.environ.get("TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS"): + raise ValueError("Collection and benchmarking require real target acceptance") + if ( + args.physical_k < 1 + or args.batch_size < 1 + or args.output_tokens < 1 + or args.warmup_cohorts < 0 + or not math.isfinite(args.timeout) + or args.timeout <= 0 + ): + raise ValueError("Positive K/BS/output/full budget and nonnegative warmups required") + if world > 1 and MPI.Query_thread() < MPI.THREAD_MULTIPLE: + raise ValueError("Concurrent ordinary worker/control MPI requires THREAD_MULTIPLE") + config = yaml.safe_load(args.config.read_text(encoding="utf-8")) + if not isinstance(config, dict) or not isinstance(config.get("speculative_config"), dict): + raise ValueError("Explicit public DSpark speculative_config required") + spec = config["speculative_config"] + if spec.get("max_draft_len") != args.physical_k: + raise ValueError("CLI physical K differs from actual configuration") + numerical = copy.deepcopy(config) + for field in _POLICY_FIELDS: + numerical["speculative_config"].pop(field, None) + if args.mode == "dynamic": + if args.physical_k == 1 or args.cost_table is None: + raise ValueError("Dynamic chain needs K>1 and a structurally valid measured cost table") + spec["enable_confidence_scheduling"] = True + spec["confidence_sps_table_path"] = str(args.cost_table.resolve(strict=True)) + if args.calibration is not None: + spec["confidence_sts_path"] = str(args.calibration.resolve(strict=True)) + else: + for field in _POLICY_FIELDS: + spec.pop(field, None) + spec["enable_confidence_scheduling"] = False + if args.cost_table is not None or args.calibration is not None: + raise ValueError("Static/calibration collection does not consume dynamic artifacts") + llm_args = TorchLlmArgs(**{**config, "model": str(args.model), "backend": "pytorch"}) + mapping = llm_args.parallel_config.to_mapping() + factor = world if mapping.enable_attention_dp else 1 + if ( + not isinstance(llm_args.speculative_config, DSparkDecodingConfig) + or not llm_args.speculative_config.draft_is_embedded_in_target + or mapping.tp_size != world + or mapping.pp_size != 1 + or mapping.cp_size != 1 + or args.batch_size % factor + or llm_args.max_batch_size < args.batch_size // factor + or llm_args.cuda_graph_config is None + or not llm_args.cuda_graph_config.enable_padding + ): + raise ValueError( + "K/topology/BS capacity and padded decode graphs must match the request cohort" + ) + if ( + llm_args.guided_decoding_backend is not None + or llm_args.batched_logits_processor is not None + or llm_args.lora_config is not None + or llm_args.cache_transceiver_config is not None + or llm_args.kv_connector_config is not None + or llm_args.enable_in_graph_sampling + or llm_args.encode_only + or llm_args.enable_resource_governor + ): + raise ValueError( + "This fixed-cohort driver supports ordinary aggregate greedy generation only" + ) + if args.mode == "collect" and llm_args.disable_overlap_scheduler: + raise ValueError( + "This source-paired collector currently uses the original overlap device-input seam" + ) + data = args.inputs.read_bytes() + inputs = [parse_json(line) for line in data.splitlines() if line.strip()][: args.batch_size] + if len(inputs) != args.batch_size or any( + not isinstance(row, dict) + or not isinstance(row.get("prompt_token_ids"), list) + or not row["prompt_token_ids"] + or any(type(token) is not int or token < 0 for token in row["prompt_token_ids"]) + for row in inputs + ): + raise ValueError("One nonempty integer-token prompt per selected request required") + for index, row in enumerate(inputs): + row["input_id"] = index + if len(row["prompt_token_ids"]) + args.output_tokens > llm_args.max_seq_len: + raise ValueError("Fixed output would exceed original max_seq_len") + origin = source_origin_from_args(args) + contract = { + "physical_k": args.physical_k, + "batch_size": args.batch_size, + "output_tokens": args.output_tokens, + "inputs_sha256": canonical_sha256(inputs), + "model_config_sha256": sha256_file(args.model / "config.json"), + "numerical_configuration_sha256": canonical_sha256(numerical), + "warmup_cohorts": args.warmup_cohorts, + "sampling": {"temperature": 0.0, "top_k": 1, "ignore_eos": True, "beam_width": 1}, + } + return llm_args, inputs, contract, origin + + +def _observed_identity(worker, args, origin): + # Finite source/native/runtime provenance. HEAD is only recorded base; + # observed module bytes include this overlay and are never relabeled old. + modules = ( + "tensorrt_llm._torch.speculative.dspark", + "tensorrt_llm._torch.speculative.dspark_verify", + "tensorrt_llm._torch.speculative.dspark_reproduction", + "tensorrt_llm._torch.pyexecutor.model_engine", + "tensorrt_llm._torch.pyexecutor.py_executor", + "tensorrt_llm._torch.pyexecutor.cuda_graph_runner", + "tensorrt_llm._torch.pyexecutor.dspark_observation_runtime", + "tensorrt_llm._torch.speculative.spec_sampler_base", + "tensorrt_llm.tools.dspark_observations", + "tensorrt_llm.llmapi.llm_args", + ) + loaded = {name: sha256_file(Path(importlib.import_module(name).__file__)) for name in modules} + if source_origin_from_args(args) != origin: + raise ValueError("Observed source changed during model construction") + for name, digest in loaded.items(): + path = args.source_root / (name.replace(".", "/") + ".py") + if sha256_file(path) != digest: + raise ValueError("Actual loaded Python bytes differ from recorded source origin") + native_paths = set(torch.ops.loaded_libraries) + for name, module in tuple(sys.modules.items()): + path = getattr(module, "__file__", None) + if ( + name.startswith(("tensorrt_llm.", "torch._C")) + and path + and Path(path).suffix in (".so", ".pyd") + ): + native_paths.add(path) + native = {str(Path(path).resolve()): sha256_file(Path(path)) for path in sorted(native_paths)} + if not native: + raise ValueError("Actual native runtime identity missing") + device = torch.cuda.get_device_properties(torch.cuda.current_device()) + mapping = worker.engine.model_engine.mapping + return { + **origin, + "source_origin_observation": source_origin_input_metadata( + getattr(args, "recorded_source_origin", None), + getattr(args, "recorded_source_origin_sha256", None), + ), + "loaded_source_sha256": loaded, + "runtime": { + "torch": str(torch.__version__), + "cuda": torch.version.cuda, + "tensorrt_llm": importlib.metadata.version("tensorrt_llm"), + }, + "local_native_libraries": native, + "local_device": {"name": device.name, "major": device.major, "minor": device.minor}, + "topology": { + "tp": mapping.tp_size, + "ep": mapping.moe_ep_size, + "attention_dp": mapping.enable_attention_dp, + "pp": mapping.pp_size, + }, + "artifact_sha256": { + name: sha256_file(path) + for name, path in (("cost", args.cost_table), ("calibration", args.calibration)) + if path + }, + } + + +def run_native(args): + with isolated_control_communicator(mpi_comm()) as control: + return _run_native(args, control) + + +def _run_native(args, control): + rank, world = control.Get_rank(), control.Get_size() + deadline = time.monotonic() + args.timeout + llm_args, inputs, contract, origin = _collective( + control, "prepare", lambda: _prepare(args, world) + ) + if len(set(control.allgather(canonical_sha256((contract, origin))))) != 1: + raise ValueError("Ranks disagree on request/configuration/source inputs") + _collective( + control, + "fresh outputs", + lambda: create_private_output_directory(args.output) if rank == 0 else None, + ) + worker, record, payload = None, None, None + succeeded, shutdown_completed = False, False + try: + + def construct(): + nonlocal worker + worker = ObservationWorker( + collect_calibration=args.mode == "collect", + engine=None, + hf_model_dir=args.model, + tokenizer=None, + llm_args=llm_args, + is_llm_executor=True, + ) + return worker + + _collective(control, "construct", construct) + identity = _collective( + control, "observed identity", lambda: _observed_identity(worker, args, origin) + ) + identity["devices"] = control.allgather(identity.pop("local_device")) + identity["rank_native_libraries"] = control.allgather( + identity.pop("local_native_libraries") + ) + # Only measured runtime/source fields determine matched identity; + # artifact origin metadata is preserved separately, not rewritten. + contract["observed_runtime_sha256"] = canonical_sha256( + {key: value for key, value in identity.items() if key != "artifact_sha256"} + ) + for cohort in range(args.warmup_cohorts + 1): + _collective( + control, + "reset prior cohort routing", + lambda: worker.engine.observation_receipt(reset=True, control=control), + ) + control.Barrier() + started = time.monotonic() + responses, error = None, None + if rank == 0: + try: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("Full original startup/model budget exhausted") + responses = response_rows(worker, inputs, args.output_tokens, remaining) + except Exception as exception: + error = type(exception).__name__ + ": " + str(exception) + error, responses, elapsed = control.bcast( + (error, responses, time.monotonic() - started), root=0 + ) + if error: + raise RuntimeError(error) + _collective( + control, + "responses", + lambda: validate_responses(responses, inputs, args.output_tokens), + ) + receipt = _collective( + control, + "original drain", + lambda: worker.engine.observation_receipt(control=control), + ) + if cohort == args.warmup_cohorts: + routing = control.allgather(receipt["routing"]) + combined = { + name: sum(row[name] for row in routing) + for name in ("native_replays", "compact_replays") + } + record = { + "schema": "dspark-public-run-v1", + "mode": args.mode, + "calibration_collection": args.mode == "collect", + "inputs": inputs, + "responses": responses, + "elapsed_seconds": elapsed, + "comparison_contract": contract, + "identity": identity, + "resolved_confidence_enabled": receipt["resolved_confidence_enabled"], + "runtime_confidence_enabled": receipt["runtime_confidence_enabled"], + "routing": combined, + "rank_graph_keys": control.allgather(receipt["graph_keys"]), + "task_accuracy_qualified": False, + "production_admission_qualified": False, + } + if args.mode == "collect": + record["rank_drain"] = control.allgather(receipt["drain"]) + if args.mode == "collect": + payload = _collective( + control, "paired full-K payload", worker.engine._native_pairs.payload + ) + payload["meta"]["identity"] = identity + succeeded = True + finally: + failure = sys.exc_info()[1] + try: + _collective(control, "original shutdown", lambda: worker.shutdown() if worker else None) + shutdown_completed = worker is not None + finally: + if record is not None: + record["run_complete"] = succeeded and shutdown_completed + record["model_shutdown_completed"] = shutdown_completed + _collective( + control, + "preserve record", + lambda: write_new_json(args.output / f"rank-{rank}-run.json", record), + ) + if succeeded and shutdown_completed and payload is not None: + payload["meta"]["model_shutdown_completed"] = True + + def preserve_pairs(): + with (args.output / f"rank-{rank}-pairs.pt").open("xb") as handle: + torch.save(payload, handle) + + _collective(control, "preserve pairs", preserve_pairs) + if not succeeded or not shutdown_completed: + _collective( + control, + "preserve failure", + lambda: write_new_json( + args.output / f"rank-{rank}-failure.json", + { + "schema": "dspark-public-failure-v1", + "run_complete": False, + "failure_type": type(failure).__name__ + if failure + else "ShutdownFailure", + "model_shutdown_completed": shutdown_completed, + }, + ), + ) diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index cf7ec9d4904b..5c819ec69042 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -158,6 +158,13 @@ def _get_comm_ranks_device_id(self): return comm_ranks, device_ids + def _get_pytorch_executor_factory(self): + """Keep normal construction unchanged; offline subclasses may supply a factory.""" + from tensorrt_llm._torch.pyexecutor.py_executor_creator import \ + create_py_executor + + return create_py_executor + def setup_engine(self): """ Setup the engine for the worker. @@ -173,9 +180,7 @@ def _create_py_executor(): ), "llm_args should be with backend in _create_py_executor" _ = self._get_comm_ranks_device_id() if self._backend == "pytorch": - from tensorrt_llm._torch.pyexecutor.py_executor_creator import \ - create_py_executor - create_executor = create_py_executor + create_executor = self._get_pytorch_executor_factory() args["llm_args"] = self.llm_args args["checkpoint_dir"] = self._hf_model_dir args["tokenizer"] = self._tokenizer diff --git a/tensorrt_llm/tools/dspark_calibration.py b/tensorrt_llm/tools/dspark_calibration.py new file mode 100644 index 000000000000..e705922cb234 --- /dev/null +++ b/tensorrt_llm/tools/dspark_calibration.py @@ -0,0 +1,239 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Offline chain-confidence STS fitting from correctly paired observation shards. + +This CPU tool imports PyTorch, not TensorRT-LLM or CUDA. Collection is a separate +runtime adapter: each row must join raw logits to full-block prefix-accept labels +from the same stamped draft sequence. Calibration is not task qualification or +an economic admission receipt. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +from pathlib import Path + +import torch + +__all__ = [ + "default_temperature_grid", + "expected_calibration_error", + "fit_sts_temperatures", + "load_shards", +] + +_EPS_PROB = 1e-8 + + +def default_temperature_grid() -> torch.Tensor: + """Return the established 41-point log-spaced STS grid over [0.1, 10].""" + return torch.logspace(math.log10(0.1), math.log10(10.0), steps=41) + + +def expected_calibration_error( + *, probs: torch.Tensor, targets: torch.Tensor, num_bins: int +) -> float: + """Compute population-weighted cumulative-prefix calibration error.""" + probs = probs.reshape(-1).to(torch.float64).clamp(_EPS_PROB, 1.0 - _EPS_PROB) + targets = targets.reshape(-1).to(torch.float64) + total = probs.numel() + if total == 0: + return float("nan") + bin_index = (probs * num_bins).long().clamp_(0, num_bins - 1) + count = torch.zeros(num_bins, dtype=torch.float64) + pred_sum = torch.zeros(num_bins, dtype=torch.float64) + target_sum = torch.zeros(num_bins, dtype=torch.float64) + count.scatter_add_(0, bin_index, torch.ones_like(probs)) + pred_sum.scatter_add_(0, bin_index, probs) + target_sum.scatter_add_(0, bin_index, targets) + denom = count.clamp_min(1.0) + bin_error = (pred_sum / denom - target_sum / denom).abs() + return float((bin_error * count).sum().item() / total) + + +def _validate_rows(logits: torch.Tensor, prefix_mask: torch.Tensor) -> None: + if ( + not isinstance(logits, torch.Tensor) + or not isinstance(prefix_mask, torch.Tensor) + or logits.device.type != "cpu" + or prefix_mask.device.type != "cpu" + or not logits.is_floating_point() + or logits.ndim != 2 + or logits.shape != prefix_mask.shape + or logits.shape[0] == 0 + or logits.shape[1] == 0 + ): + raise ValueError("CPU logits and prefix_mask must have identical nonempty [N, K] shapes") + if not torch.isfinite(logits).all() or not torch.isfinite(prefix_mask).all(): + raise ValueError("Calibration observations must be finite") + if not ((prefix_mask == 0) | (prefix_mask == 1)).all(): + raise ValueError("Prefix labels must be binary full-block acceptance observations") + if (prefix_mask[:, 1:] > prefix_mask[:, :-1]).any(): + raise ValueError("A later accepted prefix cannot follow an earlier rejected prefix") + + +def fit_sts_temperatures( + *, logits: torch.Tensor, prefix_mask: torch.Tensor, grid: torch.Tensor, num_bins: int = 15 +) -> dict[str, list[float]]: + """Fit each temperature left-to-right against cumulative survival ECE. + + Args: + logits: Raw confidence logits, CPU tensor [N, K]. + prefix_mask: Binary cumulative-prefix labels, CPU tensor [N, K]. + grid: Positive finite temperature candidates, CPU tensor [M]. + num_bins: Positive number of calibration-error bins. + """ + _validate_rows(logits, prefix_mask) + if type(num_bins) is not int or num_bins < 1: + raise ValueError("num_bins must be a positive integer") + if ( + not isinstance(grid, torch.Tensor) + or grid.device.type != "cpu" + or grid.ndim != 1 + or grid.numel() == 0 + or not torch.isfinite(grid).all() + or not (grid > 0).all() + ): + raise ValueError("The temperature grid must be a nonempty positive finite CPU vector") + logits = logits.to(torch.float64) + prefix_mask = prefix_mask.to(torch.float64) + num_samples, block_size = logits.shape + grid_values = grid.to(torch.float64).tolist() + temperatures, ece_before, ece_after = [], [], [] + survival_at_one = torch.ones(num_samples, dtype=torch.float64) + survival_fitted = torch.ones(num_samples, dtype=torch.float64) + for position in range(block_size): + position_logits = logits[:, position] + position_target = prefix_mask[:, position] + survival_at_one = survival_at_one * torch.sigmoid(position_logits) + ece_before.append( + expected_calibration_error( + probs=survival_at_one, targets=position_target, num_bins=num_bins + ) + ) + best_t = grid_values[0] + best_survival = survival_fitted * torch.sigmoid(position_logits / best_t) + best_ece = expected_calibration_error( + probs=best_survival, targets=position_target, num_bins=num_bins + ) + for temperature in grid_values[1:]: + candidate = survival_fitted * torch.sigmoid(position_logits / temperature) + ece = expected_calibration_error( + probs=candidate, targets=position_target, num_bins=num_bins + ) + if ece < best_ece: + best_ece, best_t, best_survival = ece, temperature, candidate + temperatures.append(float(best_t)) + ece_after.append(float(best_ece)) + survival_fitted = best_survival + return {"temperatures": temperatures, "ece_before": ece_before, "ece_after": ece_after} + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while block := stream.read(1 << 20): + digest.update(block) + return digest.hexdigest() + + +def load_shards( + paths: list[Path], physical_k: int +) -> tuple[torch.Tensor, torch.Tensor, dict[str, str]]: + """Load safe CPU tensor shards without conflating widths or label pairings. + + Pairing metadata is a producer contract, not independent proof of the + running checkpoint or label provenance. The original shard hashes remain + recorded; source or GPU freshness is not rewritten or certified. + """ + if type(physical_k) is not int or physical_k < 1 or not paths: + raise ValueError("An explicit positive physical K and input shards are required") + canonical = [path.resolve(strict=True) for path in paths] + if len(set(canonical)) != len(canonical): + raise ValueError("Duplicate calibration shards would double-count samples") + logits, masks, hashes = [], [], {} + for path, resolved in zip(paths, canonical): + if path.is_symlink() or not resolved.is_file(): + raise ValueError("Calibration inputs must be regular, nonsymlink files") + before = resolved.stat() + pin = _sha256_file(resolved) + blob = torch.load(resolved, map_location="cpu", weights_only=True) + if not isinstance(blob, dict) or not isinstance(blob.get("meta"), dict): + raise ValueError("Calibration shards require explicit pairing metadata") + if blob["meta"].get("pairing") != "draft_seq_ring": + raise ValueError("Recollect labels paired with the same stamped draft sequence") + raw_logits, prefix_mask = blob.get("logits"), blob.get("prefix_mask") + _validate_rows(raw_logits, prefix_mask) + if raw_logits.shape[1] != physical_k: + raise ValueError("The calibration shard width does not match the explicit physical K") + after = resolved.stat() + if (before.st_dev, before.st_ino, before.st_size, before.st_mtime_ns) != ( + after.st_dev, + after.st_ino, + after.st_size, + after.st_mtime_ns, + ) or _sha256_file(resolved) != pin: + raise ValueError("Calibration shard changed during reading") + logits.append(raw_logits) + masks.append(prefix_mask) + hashes[str(path)] = pin + return torch.cat(logits, 0), torch.cat(masks, 0), hashes + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, allow_abbrev=False) + parser.add_argument("--shards", nargs="+", type=Path, required=True) + parser.add_argument("--physical-k", type=int, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--num-bins", type=int, default=15) + parser.add_argument("--minimum-samples", type=int, default=1000) + args = parser.parse_args(argv) + if args.minimum_samples < 1: + raise ValueError("minimum-samples must be positive") + logits, masks, hashes = load_shards(args.shards, args.physical_k) + if logits.shape[0] < args.minimum_samples: + raise ValueError("Not enough paired observations to fit the requested calibration") + result = fit_sts_temperatures( + logits=logits, prefix_mask=masks, grid=default_temperature_grid(), num_bins=args.num_bins + ) + rates = masks.to(torch.float64).mean(dim=0).tolist() + payload = { + "sts_temperatures": result["temperatures"], + **result, + "physical_k": args.physical_k, + "num_samples": int(logits.shape[0]), + "num_bins": args.num_bins, + "raw_shard_sha256": hashes, + "pairing": "draft_seq_ring", + "prefix_acceptance_rate": rates, + "degenerate_positions": [index for index, rate in enumerate(rates) if rate in (0.0, 1.0)], + "objective": "left-to-right cumulative-prefix ECE;41-point temperature grid", + "producer_pairing_independently_verified": False, + "task_accuracy_or_economic_admission_qualified": False, + "artifact_freshness_validated": False, + } + if args.output.is_symlink(): + raise ValueError("Calibration output cannot be a symlink") + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0) + with os.fdopen(os.open(args.output, flags, 0o600), "w", encoding="utf-8") as stream: + json.dump(payload, stream, sort_keys=True, allow_nan=False) + stream.write("\n") + print( + json.dumps( + { + "output": str(args.output), + "num_samples": payload["num_samples"], + "physical_k": args.physical_k, + "degenerate_positions": payload["degenerate_positions"], + } + ) + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tensorrt_llm/tools/dspark_cost_export.py b/tensorrt_llm/tools/dspark_cost_export.py new file mode 100644 index 000000000000..b14e153ae409 --- /dev/null +++ b/tensorrt_llm/tools/dspark_cost_export.py @@ -0,0 +1,413 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Export direct DSpark whole-step costs or check existing artifacts. + +This standalone command uses the actual schema-v2 consumer. Only the explicit +complete-loop producer can export new costs; graph replay timings stay separate. +Run it by file path to avoid importing the native TensorRT-LLM package on CPU. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.util +import json +import os +import re +import statistics +import subprocess +import sys +from pathlib import Path +from types import ModuleType + +__all__ = [ + "check_cost_bytes", + "main", + "observed_engine_fingerprint", + "recorded_source_origin", + "source_origin_from_args", + "source_origin_input_metadata", + "runtime_cost_bytes", +] + +WHOLE_STEP_SCOPE = "original-executor-loop-hostStepTimeMS" + + +def _canonical_sha256(value: object) -> str: + return hashlib.sha256( + json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() + ).hexdigest() + + +def source_origin_input_metadata( + recorded_origin: Path | None, recorded_origin_sha256: str | None +) -> dict[str, str | bool]: + """Label recorded provenance separately from a fresh Git observation.""" + if (recorded_origin is None) != (recorded_origin_sha256 is None): + raise ValueError("Recorded source origin requires both its path and SHA256") + if recorded_origin is None: + return {"method": "git", "recorded_input": False} + if not isinstance(recorded_origin_sha256, str) or not re.fullmatch( + "[0-9a-f]{64}", recorded_origin_sha256 + ): + raise ValueError("Recorded source origin SHA256 must be explicit lowercase hex") + return { + "method": "recorded_provenance_input", + "recorded_input": True, + "input_sha256": recorded_origin_sha256, + "fresh_git_diff_repeated": False, + } + + +def source_origin_from_args(args: argparse.Namespace) -> dict[str, str]: + """Use one shared optional provenance input across the public drivers.""" + return recorded_source_origin( + args.source_root, + recorded_origin=getattr(args, "recorded_source_origin", None), + recorded_origin_sha256=getattr(args, "recorded_source_origin_sha256", None), + ) + + +def recorded_source_origin( + source_root: Path, + *, + recorded_origin: Path | None = None, + recorded_origin_sha256: str | None = None, +) -> dict[str, str]: + """Read Git provenance or an explicitly pinned record for this exact checkout. + + An optional record contains actual earlier HEAD/full-diff observation for + the same immutable source root. It is recorded metadata, not proof of source + freshness, GPU health, numerical accuracy or performance admission. Loaded + module checks and caller-owned source byte validation remain independent. + """ + source_origin_input_metadata(recorded_origin, recorded_origin_sha256) + if recorded_origin is not None: + if recorded_origin.is_symlink() or not recorded_origin.is_file(): + raise ValueError("Recorded source origin must be a regular non-symlink file") + if recorded_origin.stat().st_size > 8192: + raise ValueError("Recorded source origin must be a finite metadata record") + data = recorded_origin.read_bytes() + if hashlib.sha256(data).hexdigest() != recorded_origin_sha256: + raise ValueError("Recorded source origin bytes differ from the supplied SHA256") + utility = _source_module( + "_dspark_recorded_origin_json", Path(__file__).with_name("verify_costs.py") + ) + record = utility.parse_json(data) + if ( + not isinstance(record, dict) + or set(record) != {"schema", "source_root", "origin"} + or record["schema"] != "trtllm-recorded-source-origin-v1" + or record["source_root"] != str(source_root.resolve(strict=True)) + or not source_root.is_dir() + ): + raise ValueError("Recorded source origin belongs to a different checkout or schema") + origin = record["origin"] + if ( + not isinstance(origin, dict) + or set(origin) != {"source_head", "source_diff_sha256"} + or not isinstance(origin["source_head"], str) + or not re.fullmatch("[0-9a-f]{40}", origin["source_head"]) + or not isinstance(origin["source_diff_sha256"], str) + or not re.fullmatch("[0-9a-f]{64}", origin["source_diff_sha256"]) + ): + raise ValueError("Recorded source HEAD/full-diff digests are malformed") + return origin + head = ( + subprocess.run( + ["git", "--no-optional-locks", "-C", str(source_root), "rev-parse", "HEAD"], + check=True, + capture_output=True, + timeout=30, + ) + .stdout.decode() + .strip() + ) + diff = subprocess.run( + [ + "git", + "--no-optional-locks", + "-C", + str(source_root), + "diff", + "--no-ext-diff", + "--no-textconv", + "--binary", + "HEAD", + "--", + ], + check=True, + capture_output=True, + timeout=30, + ).stdout + return {"source_head": head, "source_diff_sha256": hashlib.sha256(diff).hexdigest()} + + +def observed_engine_fingerprint( + identity: dict[str, object], + *, + source_root: Path, + world_size: int, + recorded_origin: Path | None = None, + recorded_origin_sha256: str | None = None, +) -> dict[str, object]: + """Record actual origin metadata, not a destination freshness certificate. + + Loaded module bytes must be those of the supplied public checkout. Native + and runtime snapshots preserve observed library/configuration identity; no + old artifact fingerprint is relabeled as current and no GPU health is implied. + """ + topology = identity["topology"] + if world_size > 1 and not topology["attention_dp"]: + raise ValueError( + "Schema-v2 global graph counts use attention-DP fanout; " + "multi-rank attention-DP-off costs need a separate schema decision" + ) + sources = identity["loaded_source_sha256"] + for module, expected in sources.items(): + path = source_root / (module.replace(".", "/") + ".py") + if not path.is_file() or hashlib.sha256(path.read_bytes()).hexdigest() != expected: + raise ValueError("Loaded public source differs from the recorded origin checkout") + devices = identity["devices"] + native = identity["rank_native_libraries"] + if len(devices) != world_size or len(native) != world_size or not all(native): + raise ValueError("Complete observed device/native-library rank coverage is required") + names = {device["name"] for device in devices} + if len(names) != 1: + raise ValueError("A cost grid must record one observed GPU class") + origin = recorded_source_origin( + source_root, recorded_origin=recorded_origin, recorded_origin_sha256=recorded_origin_sha256 + ) + if identity.get("recorded_source_origin", origin) != origin: + raise ValueError("Original recorded source changed during whole-step construction") + graphs = sorted({cell["graph_batch_size"] for cell in identity["profile"]["cells"]}) + runtime = { + "runtime": identity["runtime"], + "rank_native_libraries": native, + "loaded_source_sha256": sources, + "runtime_configuration": identity["runtime_configuration"], + "model_config_sha256": identity["model_config_sha256"], + "profile": identity["profile"], + } + return { + "global_graph_batch_sizes": [g * world_size for g in graphs], + "gpu": next(iter(names)), + "gpu_count": world_size, + "gpu_snapshot_sha256": _canonical_sha256(devices), + "max_draft_len": identity["profile"]["physical_k"], + "rank_local_graph_batch_sizes": graphs, + "runtime_snapshot": _canonical_sha256(runtime), + **origin, + "topology": json.dumps(topology, sort_keys=True, separators=(",", ":")), + } + + +def _source_module(name: str, path: Path) -> ModuleType: + """Load a fixed pure-Python dependency without native package initialization.""" + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise ImportError(f"Cannot load cost tooling dependency: {path.name}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def _whole_step_payload( + data: bytes, summary: dict, *, minimum_predicted_gain: float +) -> dict[str, object]: + """Map direct complete-loop cells without adding components or missing prices.""" + if summary.get("timing_scope") != WHOLE_STEP_SCOPE: + raise ValueError( + "Graph replay observations are not schema-v2 whole-step costs; " + "use the original complete-loop whole-step producer, not relabeled timings" + ) + identity = summary["identity"] + fingerprint = identity["recorded_engine_fingerprint"] + profile = identity["profile"] + if profile["measurement_scope"] != "whole-step": + raise ValueError("Whole-step export requires the original whole-step adapter") + floors = summary["sample_floors"] + if floors["per_generation"] < 12 or floors["pooled_per_shape"] < 24: + raise ValueError("Original whole-step sample floors are 12 per generation and 24 pooled") + tables, measurements, observed = {}, [], set() + digest = hashlib.sha256(data).hexdigest() + for cell in summary["cells"]: + shape = cell["execution_shape"] + g, v = shape["graph_batch_size"], shape["verifier_budget"] + expected_layout = "uniform" if v == 0 else "packed_chain_rows" + if ( + shape["workload"] != "dspark_chain" + or shape["physical_k"] != profile["physical_k"] + or shape["attention_layout"] != expected_layout + or shape["captured_head"] != profile["record_confidence_head"] + or (g, v) in observed + ): + raise ValueError( + "Whole-step execution shape, K, layout or unique cell identity differs" + ) + samples = cell["retained_rank_max_samples_ms"] + if ( + len(samples) < floors["pooled_per_shape"] + or cell["retained_count"] != len(samples) + or any( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not 0 < value < float("inf") + for value in samples + ) + or cell["rank_max_median_ms"] != statistics.median(samples) + ): + raise ValueError("Whole-step median does not reproduce from original retained samples") + observed.add((g, v)) + table = tables.setdefault(str(g), {"token_counts": [], "step_time_ms": []}) + table["token_counts"].append(v) + table["step_time_ms"].append(cell["rank_max_median_ms"]) + measurements.append( + { + "rank_local_graph_batch_size": g, + "rank_local_verifier_budget": v, + "source_result_sha256": digest, + "step_time_ms": cell["rank_max_median_ms"], + } + ) + declared = {(cell["graph_batch_size"], cell["verifier_budget"]) for cell in profile["cells"]} + if observed != declared: + raise ValueError( + "Every declared native/compact whole-step cell requires its direct measurement" + ) + for table in tables.values(): + ordered = sorted(zip(table["token_counts"], table["step_time_ms"])) + table["token_counts"], table["step_time_ms"] = [list(values) for values in zip(*ordered)] + return { + "schema_version": 2, + "cost_tables": tables, + "engine_fingerprint": fingerprint, + "engine_fingerprint_sha256": _canonical_sha256(fingerprint), + "minimum_predicted_gain": minimum_predicted_gain, + "measurements": sorted( + measurements, + key=lambda item: ( + item["rank_local_graph_batch_size"], + item["rank_local_verifier_budget"], + ), + ), + } + + +def runtime_cost_bytes( + data: bytes, + *, + physical_k: int, + graph_batch_sizes: list[int], + minimum_predicted_gain: float = 0.01, +) -> tuple[bytes, dict[str, object]]: + """Validate reuse or serialize measured complete-loop costs into existing schema. + + A raw export summary has a different timing/identity contract. Even if a + caller toggles an eligibility flag, that cannot extend its measured scope. + Already supplied schema-v2 artifacts retain their original provenance and + values; validation is not a new GPU measurement or admission decision. + """ + utilities = _source_module( + "_dspark_cost_export_utilities", Path(__file__).with_name("verify_costs.py") + ) + payload = utilities.parse_json(data) + if not isinstance(payload, dict): + raise TypeError("Costs must contain a JSON object") + whole_step = payload.get("schema") == "verify-cost-measured-summary-v1" + if whole_step: + payload = _whole_step_payload(data, payload, minimum_predicted_gain=minimum_predicted_gain) + elif payload.get("schema") == "verify-cost-raw-v1": + raise ValueError( + "Raw observations first require exact-rank reduction from the whole-step producer" + ) + consumer_path = Path(__file__).resolve().parents[1] / "_torch/speculative/dspark_planner.py" + consumer = _source_module("_dspark_cost_export_consumer", consumer_path) + fingerprint = consumer.validate_sps_cost_table_payload( + payload, + graph_batch_sizes=graph_batch_sizes, + max_draft_len=physical_k, + validate_freshness=False, + ) + # Use the same constructor as the file loader, including its bounded exact + # candidate set and any already supplied drain metadata. No policy is added. + consumer._build_exact_sps_cost_table( + payload, max_draft_len=physical_k, validate_artifact_freshness=False + ) + output = ( + json.dumps(payload, sort_keys=True, allow_nan=False).encode() + b"\n" + if whole_step + else data + ) + return output, { + "schema": "dspark-existing-cost-structural-check-v1", + "source_artifact_sha256": hashlib.sha256(data).hexdigest(), + "runtime_cost_artifact_sha256": hashlib.sha256(output).hexdigest(), + "measured_timing_scope": WHOLE_STEP_SCOPE if whole_step else "original-supplied-artifact", + "recorded_engine_fingerprint_sha256": payload["engine_fingerprint_sha256"], + "recorded_source_head": fingerprint["source_head"], + "physical_k": physical_k, + "rank_local_graph_batch_sizes": graph_batch_sizes, + "structural_validation_passed": True, + "destination_freshness_checked": False, + "new_measurements_created": False, + "direct_whole_step_measurements_exported": whole_step, + "task_accuracy_qualified": False, + "performance_or_economic_admission_qualified": False, + "warning": ( + "Supplied predictions retain their original source and GPU identity; " + "their freshness and performance on the destination are unchecked." + ), + } + + +def check_cost_bytes( + data: bytes, *, physical_k: int, graph_batch_sizes: list[int] +) -> dict[str, object]: + """Return the existing artifact or complete-loop export's structural report.""" + return runtime_cost_bytes(data, physical_k=physical_k, graph_batch_sizes=graph_batch_sizes)[1] + + +def _write_new_bytes(path: Path, data: bytes) -> None: + """Preserve every original byte; never overwrite or follow a final symlink.""" + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0) + with os.fdopen(os.open(path, flags, 0o600), "wb") as stream: + stream.write(data) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input", type=Path, required=True) + parser.add_argument("--physical-k", type=int, required=True) + parser.add_argument("--graph-batch-sizes", type=int, nargs="+", required=True) + parser.add_argument( + "--minimum-predicted-gain", + type=float, + default=0.01, + help="Policy margin for a newly exported whole-step table only.", + ) + parser.add_argument( + "--output", type=Path, help="Optional new copy; original bytes are preserved." + ) + args = parser.parse_args(argv) + if args.input.is_symlink() or not args.input.is_file(): + raise ValueError("Input must be a regular, non-symlink cost artifact") + data = args.input.read_bytes() + output, report = runtime_cost_bytes( + data, + physical_k=args.physical_k, + graph_batch_sizes=args.graph_batch_sizes, + minimum_predicted_gain=args.minimum_predicted_gain, + ) + if args.input.read_bytes() != data: + raise ValueError("Original cost artifact changed during validation") + if args.output is not None: + _write_new_bytes(args.output, output) + print(json.dumps(report, sort_keys=True, allow_nan=False)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tensorrt_llm/tools/dspark_observations.py b/tensorrt_llm/tools/dspark_observations.py new file mode 100644 index 000000000000..d6745a4f4a20 --- /dev/null +++ b/tensorrt_llm/tools/dspark_observations.py @@ -0,0 +1,213 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU joins for offline DSpark observations; no timing or policy decisions.""" + +from __future__ import annotations + +import dataclasses +import hashlib +import json +import math +from collections import Counter +from typing import Any + +import torch + + +def token_digest(tokens: list[int]) -> str: + return hashlib.sha256(json.dumps(tokens, separators=(",", ":")).encode()).hexdigest() + + +@dataclasses.dataclass +class ProposalSnapshot: + """One genuine forward's proposals and confidence under the same sampler event. + + Host tensors must not be read until the ORIGINAL sampler has updated this + state. ``device`` ownership identifies which proposal the later forward + consumed; row order or the latest live confidence never establishes pairing. + """ + + iteration: int + state: Any + confidence_host: torch.Tensor + requests: dict[int, tuple[Any, int, int]] # original object, sampler slot, confidence row + stamp_host: torch.Tensor + draft_sequence: int + updated: bool = False + + +class NativePairs: + """Accept only uncensored, incarnation- and proposal-exact full-K pairs.""" + + def __init__(self, physical_k: int) -> None: + if type(physical_k) is not int or physical_k < 1: + raise ValueError("Positive physical K required") + self.physical_k = physical_k + self.logits: list[torch.Tensor] = [] + self.accepted: list[int] = [] + self.joins: list[dict] = [] + self._pair_keys: set[tuple[int, int, int, int]] = set() + self.stats: Counter = Counter() + + def record( + self, + *, + source: ProposalSnapshot, + request: Any, + target_iteration: int, + input_proposal: torch.Tensor, + accepted: int, + verified: int, + target_slot: int, + ) -> None: + if not source.updated: + raise RuntimeError("Original source sampler update has not completed") + if target_iteration <= source.iteration: + raise ValueError("A proposal must be verified after its producing forward") + record = source.requests.get(request.py_request_id) + if record is None or record[0] is not request: + raise ValueError("Proposal belongs to another request incarnation") + if type(accepted) is not int or not 0 <= accepted <= self.physical_k: + raise ValueError("Full-K accepted-prefix count outside physical K") + if verified != self.physical_k: + raise ValueError("Compact or partial verification censors calibration labels") + _, source_slot, confidence_row = record + if ( + source.stamp_host.device.type != "cpu" + or int(source.stamp_host[confidence_row]) != source.draft_sequence + ): + raise ValueError("Confidence row was not written by the bound producing draft") + source_tokens = source.state.host.next_draft_tokens[source_slot] + if ( + source_tokens.device.type != "cpu" + or input_proposal.device.type != "cpu" + or source_tokens.numel() != self.physical_k + or input_proposal.numel() != self.physical_k + or not torch.equal(source_tokens, input_proposal) + ): + raise ValueError("Actual verification input differs from the bound source proposal") + logits = source.confidence_host[confidence_row] + if ( + logits.device.type != "cpu" + or logits.shape != (self.physical_k,) + or not logits.is_floating_point() + or not torch.isfinite(logits).all() + ): + raise ValueError("Bound confidence must contain finite full-K logits") + key = (request.py_request_id, id(request), source.iteration, target_iteration) + if key in self._pair_keys: + raise ValueError("Repeated calibration pair") + self._pair_keys.add(key) + self.logits.append(logits.clone()) + self.accepted.append(accepted) + self.joins.append( + { + "pair_key": list(key), + "source_slot": source_slot, + "target_slot": target_slot, + "confidence_row": confidence_row, + "proposal_sha256": token_digest(source_tokens.tolist()), + "accepted_draft_prefix": accepted, + "verified_drafts": verified, + "producing_draft_sequence": source.draft_sequence, + } + ) + self.stats["recorded"] += 1 + + def payload(self) -> dict: + if not self.logits: + raise ValueError("No genuine paired full-K observations") + counts = torch.tensor(self.accepted, dtype=torch.int64).view(-1, 1) + return { + "logits": torch.stack(self.logits), + "prefix_mask": (torch.arange(self.physical_k).view(1, -1) < counts).float(), + "meta": { + "pairing": "draft_seq_ring", + "pairing_protocol": "original-input-state-object-and-proposal-v1", + "block_size": self.physical_k, + "stats": dict(self.stats), + "joins": self.joins, + "native_full_k": True, + "task_accuracy_qualified": False, + }, + } + + +def validate_responses(rows: list[dict], inputs: list[dict], output_tokens: int) -> None: + if len(rows) != len(inputs) or len({row["input_id"] for row in rows}) != len(rows): + raise ValueError("Missing, duplicated or additional original response") + expected = {row["input_id"]: row["prompt_token_ids"] for row in inputs} + for row in rows: + tokens = row["output_token_ids"] + if ( + row["input_id"] not in expected + or row["prompt_token_ids"] != expected[row["input_id"]] + or row["finish_reason"] != "length" + or len(tokens) != output_tokens + or any(type(token) is not int or token < 0 for token in tokens) + ): + raise ValueError("Original prompt, output cardinality or token contract failed") + + +def compare_runs(records: list[dict]) -> dict: + """Matched fixed-cohort E2E comparison, not task accuracy or a best-K claim.""" + arms: dict[str, list[dict]] = {"static": [], "dynamic": []} + contract = None + for record in records: + if ( + record.get("schema") != "dspark-public-run-v1" + or not record.get("run_complete") + or not record.get("model_shutdown_completed") + ): + raise ValueError("Only naturally completed, closed public runs can be compared") + mode = record["mode"] + if mode not in arms or record.get("calibration_collection"): + raise ValueError("Calibration observation is not a serving benchmark arm") + match = record["comparison_contract"] + if contract is None: + contract = match + if match != contract: + raise ValueError("Inputs, physical K, BS or numerical runtime differ") + if bool(record["resolved_confidence_enabled"]) != (mode == "dynamic") or bool( + record["runtime_confidence_enabled"] + ) != (mode == "dynamic"): + raise ValueError("Requested dynamic execution fell back to static") + routes = record["routing"] + if not routes["native_replays"] and not routes["compact_replays"]: + raise ValueError("No real generation graph replay") + if mode == "static" and routes["compact_replays"]: + raise ValueError("Static arm unexpectedly used compact verification") + if mode == "dynamic" and routes["compact_replays"] <= 0: + raise ValueError("Enabled policy never executed compact; no dynamic gain claim") + value = record["elapsed_seconds"] + if type(value) not in (int, float) or not math.isfinite(value) or value <= 0: + raise ValueError("Positive real E2E elapsed time required") + validate_responses(record["responses"], record["inputs"], contract["output_tokens"]) + arms[mode].append(record) + if not all(arms.values()) or len(arms["static"]) != len(arms["dynamic"]): + raise ValueError("Equal numbers of matched static and dynamic lifetimes required") + token_work = contract["batch_size"] * contract["output_tokens"] + prices = { + arm: token_work * len(runs) / sum(run["elapsed_seconds"] for run in runs) + for arm, runs in arms.items() + } + native_outputs = { + row["input_id"]: row["output_token_ids"] for row in arms["static"][0]["responses"] + } + equal = all( + {row["input_id"]: row["output_token_ids"] for row in run["responses"]} == native_outputs + for runs in arms.values() + for run in runs + ) + return { + "schema": "dspark-public-comparison-v1", + "comparison_contract": contract, + "timing_scope": "fixed-cohort-request-enqueue-to-final-response-including-prefill-and-drain", + "tps": prices, + "dynamic_over_static_pct": 100 * (prices["dynamic"] / prices["static"] - 1), + "exact_output_tokens_equal": equal, + "task_accuracy_qualified": False, + "default_production_gain_qualified": False, + "confidence_interval": None, + "best_static_or_dynamic_k_qualified": False, + } diff --git a/tensorrt_llm/tools/dspark_reproduction.py b/tensorrt_llm/tools/dspark_reproduction.py new file mode 100644 index 000000000000..dc318763095d --- /dev/null +++ b/tensorrt_llm/tools/dspark_reproduction.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Public native calibration collection and normal static/enabled comparison. + +Run by file path so CPU-only comparison/help does not initialize TRT-LLM/CUDA. +Each ``run`` is one independent model lifetime; order/repeats belong to caller. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import sys +from pathlib import Path + + +def _cpu_module(): + path = Path(__file__).with_name("dspark_observations.py") + name = "_dspark_public_observations" + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__, allow_abbrev=False) + sub = parser.add_subparsers(dest="command", required=True) + run = sub.add_parser("run", help="One real full-model lifetime, with natural shutdown") + run.add_argument("--mode", choices=("static", "dynamic", "collect"), required=True) + run.add_argument("--model", type=Path, required=True) + run.add_argument("--config", type=Path, required=True) + run.add_argument("--inputs", type=Path, required=True) + run.add_argument("--source-root", type=Path, required=True) + run.add_argument( + "--recorded-source-origin", + type=Path, + help="Optional recorded HEAD/full-diff JSON for this immutable checkout.", + ) + run.add_argument( + "--recorded-source-origin-sha256", + help="SHA256 of the exact optional recorded source-origin bytes.", + ) + run.add_argument("--physical-k", type=int, required=True) + run.add_argument( + "--batch-size", type=int, required=True, help="Global fixed-cohort request count" + ) + run.add_argument("--output-tokens", type=int, required=True) + run.add_argument( + "--timeout", type=float, required=True, help="Full startup+model+shutdown budget in seconds" + ) + run.add_argument("--warmup-cohorts", type=int, default=1) + run.add_argument("--cost-table", type=Path) + run.add_argument("--calibration", type=Path) + run.add_argument("--output", type=Path, required=True, help="Fresh private output directory") + compare = sub.add_parser("compare", help="Compare closed independent model lifetimes") + compare.add_argument("--runs", type=Path, nargs="+", required=True) + compare.add_argument("--output", type=Path, required=True) + args = parser.parse_args(argv) + if args.command == "run": + from tensorrt_llm._torch.speculative.dspark_reproduction import run_native + + return run_native(args) + records = [json.loads(path.read_bytes()) for path in args.runs] + report = _cpu_module().compare_runs(records) + with args.output.open("x", encoding="utf-8") as handle: + json.dump(report, handle, indent=2, allow_nan=False) + + +if __name__ == "__main__": + main() diff --git a/tensorrt_llm/tools/verify_costs.py b/tensorrt_llm/tools/verify_costs.py new file mode 100644 index 000000000000..1888a83bf41f --- /dev/null +++ b/tensorrt_llm/tools/verify_costs.py @@ -0,0 +1,445 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Workload-independent offline GPU measurements and strict raw-cost export. + +This module needs no CUDA or native TRT-LLM import to reduce existing measurements. +Its exports are not serving cost tables, accuracy certificates or admission receipts. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +import re +import statistics +from collections import defaultdict +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Callable, Protocol + +__all__ = [ + "EventLeasePool", + "canonical_sha256", + "parse_json", + "sha256_file", + "reduce_measurements", + "write_new_json", + "isolated_control_communicator", + "create_private_output_directory", +] + + +class TimingEvent(Protocol): + def record(self, stream=None) -> None: ... + def query(self) -> bool: ... + def elapsed_time(self, other: TimingEvent) -> float: ... + + +class ControlCommunicator(Protocol): + def Dup(self) -> ControlCommunicator: ... + def Free(self) -> None: ... + + +@contextmanager +def isolated_control_communicator(runtime: ControlCommunicator) -> Iterator[ControlCommunicator]: + """Duplicate collectively before workers; never modify the runtime communicator. + + The caller must finish worker shutdown and evidence closure inside this + context. Its dedicated communicator remains owned on every error path. + """ + control = runtime.Dup() + try: + yield control + finally: + control.Free() + + +class EventLeasePool: + """Keep bounded event pairs alive until completed-row readback releases them. + + The caller materializes CUDA events before capture and reads them only after + its normal completion boundary. No per-step synchronization is introduced. + """ + + def __init__(self, event_factory: Callable[[], TimingEvent], size: int = 32) -> None: + if type(size) is not int or not 2 <= size <= 256: + raise ValueError("Event pool size must be in 2..256") + self._events = [(event_factory(), event_factory()) for _ in range(size)] + self._owners: dict[int, int] = {} + self._ended: set[int] = set() + self._next_id = 0 + self._initialized = False + + def initialize(self, stream=None) -> None: + """Materialize lazy handles outside capture; caller establishes completion.""" + if self._initialized: + raise RuntimeError("Event pool already initialized") + for start, end in self._events: + start.record(stream) + end.record(stream) + self._initialized = True + + def begin(self) -> int: + if not self._initialized: + raise RuntimeError("Initialize timing events before capture") + occupied = set(self._owners.values()) + slot = next((i for i in range(len(self._events)) if i not in occupied), None) + if slot is None: + raise RuntimeError("No completed event lease available") + lease = self._next_id + self._next_id += 1 + self._owners[lease] = slot + self._events[slot][0].record() + return lease + + def end(self, lease: int) -> None: + if type(lease) is not int or lease not in self._owners: + raise RuntimeError("Unknown timing lease") + if lease in self._ended: + raise RuntimeError("Timing lease ended twice") + self._events[self._owners[lease]][1].record() + self._ended.add(lease) + + def complete(self, lease: int) -> float: + if type(lease) is not int or lease not in self._owners: + raise RuntimeError("Unknown timing lease") + if lease not in self._ended: + raise RuntimeError("Timing lease has no end boundary") + start, end = self._events[self._owners[lease]] + if not end.query(): + raise RuntimeError("Normal completion did not finish the timing events") + elapsed = float(start.elapsed_time(end)) + if not math.isfinite(elapsed) or elapsed <= 0: + raise ValueError("Measured duration must be positive and finite") + del self._owners[lease] + self._ended.remove(lease) + return elapsed + + @property + def pending(self) -> int: + return len(self._owners) + + +def canonical_sha256(value: object) -> str: + return hashlib.sha256( + json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() + ).hexdigest() + + +def parse_json(data: bytes | str) -> object: + """Reject ambiguous duplicate keys and nonstandard nonfinite JSON values.""" + + def pairs(items): + result = {} + for key, value in items: + if key in result: + raise ValueError(f"Duplicate JSON key: {key}") + result[key] = value + return result + + def constant(value): + raise ValueError(f"Nonfinite JSON value: {value}") + + return json.loads(data, object_pairs_hook=pairs, parse_constant=constant) + + +def sha256_file(path: Path) -> str: + """Hash actual artifact bytes without loading a whole native library into RAM.""" + digest = hashlib.sha256() + with path.open("rb") as stream: + while block := stream.read(1 << 20): + digest.update(block) + return digest.hexdigest() + + +def write_new_json(path: Path, value: object) -> None: + """Create a new owner-only POSIX file; never follow/overwrite an existing target. + + No existing ACL, umask or host setting is changed. POSIX permission semantics + need Linux qualification; Windows ACL behavior is not attested by this helper. + """ + payload = json.dumps(value, sort_keys=True, allow_nan=False) + "\n" + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0) + with os.fdopen(os.open(path, flags, 0o600), "w", encoding="utf-8") as stream: + stream.write(payload) + + +def create_private_output_directory(path: Path) -> None: + """Create only a new diagnostic directory, with owner-only POSIX permissions.""" + path.mkdir(mode=0o700, parents=False, exist_ok=False) + + +def reduce_measurements( + documents: list[dict], + *, + excluded_generations: int = 1, + warmup_samples: int = 2, + minimum_samples: int = 24, + minimum_samples_per_generation: int = 12, +) -> dict: + """Join exact rank/sample identities before reducing rank maxima. + + Warm-up exclusions remain in the output. Missing ranks, unequal ordinal sets, + mixed identities or scopes, and incomplete lifecycle closure are errors. + Samples are not summed across ranks or across overlapping timing scopes. + """ + for value in (excluded_generations, warmup_samples): + if type(value) is not int or value < 0: + raise ValueError("Warm-up exclusions must be nonnegative integers") + if ( + type(minimum_samples) is not int + or minimum_samples < 1 + or not documents + or type(minimum_samples_per_generation) is not int + or minimum_samples_per_generation < 1 + ): + raise ValueError("A positive sample floor and raw documents are required") + if any(not isinstance(row, dict) for row in documents): + raise ValueError("Raw measurements must be JSON objects") + first = documents[0] + world = first["world_size"] + if type(world) is not int or world < 1: + raise ValueError("Actual world size required") + identity = first["identity"] + scope = first["timing_scope"] + if not isinstance(scope, str) or not scope or not isinstance(identity, dict) or not identity: + raise ValueError("Observed identity and explicit timing scope required") + declared = identity.get("declared_shape_ids", []) + if ( + not isinstance(declared, list) + or not declared + or any( + not isinstance(key, str) or not re.fullmatch(r"[0-9a-f]{64}", key) for key in declared + ) + or len(set(declared)) != len(declared) + ): + raise ValueError("Adapter-declared finite shape identities required") + expected_generations = identity.get("generations_per_shape") + if type(expected_generations) is not int or expected_generations < 1: + raise ValueError("Declared generations per shape required") + seen_shapes = set() + declared_to_observed = {} + cells: dict[str, dict[int, dict[int, dict]]] = defaultdict(lambda: defaultdict(dict)) + shapes = {} + raw_hashes = [] + for row in documents: + if ( + row.get("schema") != "verify-cost-raw-v1" + or type(row["world_size"]) is not int + or row["world_size"] != world + or row["identity"] != identity + or row["timing_scope"] != scope + or row.get("generation_closed") is not True + or row.get("run_complete") is not True + or row.get("model_shutdown_completed") is not True + ): + raise ValueError("Raw source, timing scope, generation or shutdown closure differs") + rank, generation = row["rank"], row["generation"] + if ( + type(rank) is not int + or not 0 <= rank < world + or type(generation) is not int + or not 0 <= generation < expected_generations + ): + raise ValueError("Invalid rank/generation identity") + if not isinstance(row["execution_shape"], dict) or not row["execution_shape"]: + raise ValueError("Explicit execution shape/layout required") + key = canonical_sha256(row["execution_shape"]) + if row.get("declared_shape_id") not in declared: + raise ValueError("Measurement is not part of the declared workload grid") + seen_shapes.add(row["declared_shape_id"]) + observed = declared_to_observed.setdefault(row["declared_shape_id"], key) + if observed != key: + raise ValueError("A declared shape changed its observed execution layout") + shapes[key] = row["execution_shape"] + if rank in cells[key][generation]: + raise ValueError("Duplicate cell/generation/rank") + cells[key][generation][rank] = row + raw_hashes.append(canonical_sha256(row)) + if seen_shapes != set(declared): + raise ValueError("A whole declared workload shape is missing") + results = [] + rank_ordinals = [dict() for _ in range(world)] + for key, generations in sorted(cells.items()): + if sorted(generations) != list(range(expected_generations)): + raise ValueError("Declared per-cell generation coverage is incomplete") + retained, excluded, joined_rows = [], [], [] + previous_maximum = -1 + for generation, ranks in sorted(generations.items()): + if set(ranks) != set(range(world)): + raise ValueError("Missing rank in a generation") + samples = [] + for rank in range(world): + rank_samples = {} + if not isinstance(ranks[rank]["samples"], list) or not ranks[rank]["samples"]: + raise ValueError("A generation has no raw samples") + for sample in ranks[rank]["samples"]: + ordinal, iteration = sample["sample_index"], sample["executor_iteration"] + duration = sample["elapsed_ms"] + if ( + type(ordinal) is not int + or ordinal < 0 + or ordinal in rank_samples + or type(iteration) is not int + or iteration < 0 + or type(duration) not in (int, float) + or not math.isfinite(duration) + or duration <= 0 + or type(sample["eligible"]) is not bool + ): + raise ValueError("Invalid or duplicated raw sample") + rank_samples[ordinal] = sample + if ordinal in rank_ordinals[rank]: + raise ValueError("A replay ordinal was reused in the run") + rank_ordinals[rank][ordinal] = iteration + samples.append(rank_samples) + ordinals = set(samples[0]) + if any(set(rank_samples) != ordinals for rank_samples in samples): + raise ValueError("Ranks have different replay ordinal sets") + if min(ordinals) <= previous_maximum: + raise ValueError("Replay ordinals reset between generations") + previous_maximum = max(ordinals) + eligible_index = 0 + generation_retained = 0 + for ordinal in sorted(ordinals): + joined = [samples[rank][ordinal] for rank in range(world)] + if len({sample["executor_iteration"] for sample in joined}) != 1: + raise ValueError("Replay ordinal joins different executor iterations") + duration = max(sample["elapsed_ms"] for sample in joined) + reason = None + if not all(sample["eligible"] for sample in joined): + reason = "non_plateau_or_incomplete_update" + elif generation < excluded_generations: + reason = "warmup_generation" + else: + if eligible_index < warmup_samples: + reason = "initial_plateau_sample" + eligible_index += 1 + detail = { + "generation": generation, + "sample_index": ordinal, + "executor_iteration": joined[0]["executor_iteration"], + "rank_elapsed_ms": [sample["elapsed_ms"] for sample in joined], + "rank_max_elapsed_ms": duration, + "excluded_reason": reason, + } + joined_rows.append(detail) + (excluded if reason else retained).append(duration) + if reason is None: + generation_retained += 1 + if ( + generation >= excluded_generations + and generation_retained < minimum_samples_per_generation + ): + raise ValueError("Insufficient retained exact samples in a generation") + if len(retained) < minimum_samples: + raise ValueError("Insufficient exact retained samples for a cell") + results.append( + { + "execution_shape": shapes[key], + "retained_count": len(retained), + "rank_max_median_ms": statistics.median(retained), + "retained_rank_max_samples_ms": retained, + "joins": joined_rows, + "excluded_count": len(excluded), + } + ) + for ordinals in rank_ordinals: + ordered = sorted(ordinals) + if ordered != list(range(len(ordered))): + raise ValueError("Whole-run replay ordinal coverage is incomplete") + iterations = [ordinals[ordinal] for ordinal in ordered] + if any(right <= left for left, right in zip(iterations, iterations[1:])): + raise ValueError("Executor iterations do not advance with replay ordinals") + return { + "schema": "verify-cost-measured-summary-v1", + "identity": identity, + "timing_scope": scope, + "cells": results, + "raw_document_sha256": raw_hashes, + "exclusions": {"generations": excluded_generations, "initial_samples": warmup_samples}, + "sample_floors": { + "per_generation": minimum_samples_per_generation, + "pooled_per_shape": minimum_samples, + }, + "production_sps_table_eligible": False, + "production_admission_qualified": False, + "policy_overhead_measured": False, + "task_accuracy_qualified": False, + } + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest="command", required=True) + collect = commands.add_parser("profile", help="Collect authentic DSpark graph replays") + collect.add_argument("--model", type=Path, required=True) + collect.add_argument("--config", type=Path, required=True) + collect.add_argument("--grid", type=Path, required=True) + collect.add_argument("--physical-k", type=int, required=True) + collect.add_argument( + "--cost-scope", choices=("graph-replay", "whole-step"), default="graph-replay" + ) + collect.add_argument( + "--source-root", + type=Path, + help="Original public checkout used for recorded whole-step metadata.", + ) + collect.add_argument("--inputs", type=Path, required=True, help="JSONL prompt_token_ids") + collect.add_argument( + "--recorded-source-origin", + type=Path, + help="Optional recorded HEAD/full-diff JSON for this immutable checkout.", + ) + collect.add_argument( + "--recorded-source-origin-sha256", + help="SHA256 of the exact optional recorded source-origin bytes.", + ) + collect.add_argument("--output", type=Path, required=True) + collect.add_argument("--generations", type=int, default=3) + collect.add_argument("--output-tokens", type=int, default=768) + collect.add_argument("--timeout", type=float, default=12180.0) + export = commands.add_parser("export", help="Reduce raw observations; never grant admission") + export.add_argument("--raw", type=Path, nargs="+", required=True) + export.add_argument("--output", type=Path, required=True) + export.add_argument("--exclude-generations", type=int, default=1) + export.add_argument("--warmup-samples", type=int, default=2) + export.add_argument("--minimum-samples", type=int, default=24) + export.add_argument("--minimum-generation-samples", type=int, default=12) + args = parser.parse_args(argv) + if ( + args.command == "profile" + and args.cost_scope != "whole-step" + and ( + args.recorded_source_origin is not None + or args.recorded_source_origin_sha256 is not None + ) + ): + parser.error("Recorded source-origin input applies only to --cost-scope whole-step") + if args.command == "profile": + from tensorrt_llm._torch.speculative.dspark_cost_profiler import run_profile + + run_profile(args) + else: + raw_sources = [(path, path.read_bytes()) for path in args.raw] + documents = [parse_json(raw) for _, raw in raw_sources] + summary = reduce_measurements( + documents, + excluded_generations=args.exclude_generations, + warmup_samples=args.warmup_samples, + minimum_samples=args.minimum_samples, + minimum_samples_per_generation=args.minimum_generation_samples, + ) + summary["raw_file_sha256"] = { + str(path): hashlib.sha256(raw).hexdigest() for path, raw in raw_sources + } + write_new_json(args.output, summary) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 54fa5dd947e2..13b447516324 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -151,6 +151,7 @@ l0_cpu: - unittest/executor/test_stop_word_side_output_trim.py - unittest/executor/test_ray_stub.py - unittest/executor/test_generation_request_array_prompt.py + - unittest/executor/test_dspark_reproduction_factory.py - unittest/grpc/openengine/test_coordinator.py - unittest/grpc/openengine/test_frontend_lifecycle.py - unittest/grpc/openengine/test_control.py diff --git a/tests/unittest/executor/test_dspark_reproduction_factory.py b/tests/unittest/executor/test_dspark_reproduction_factory.py new file mode 100644 index 000000000000..79320e399e45 --- /dev/null +++ b/tests/unittest/executor/test_dspark_reproduction_factory.py @@ -0,0 +1,397 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run in the installed native runtime on CPU before a public GPU caller. + +Actual TorchLlmArgs/LLM validation and factory signatures are exercised; only +costly model construction is intercepted. These do not prove native GPU math. +""" + +import importlib.util +import inspect +import json +import subprocess +from pathlib import Path +from types import SimpleNamespace + +import pytest +import yaml + +from tensorrt_llm._torch.pyexecutor import ( + dspark_observation_runtime, + dspark_request_profile, + py_executor_creator, +) +from tensorrt_llm._torch.speculative import dspark_reproduction +from tensorrt_llm.executor.base_worker import BaseWorker +from tensorrt_llm.llmapi.llm import LLM +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("mode", ["static", "dynamic", "collect"]) +def test_complete_public_parser_actual_llm_validator_and_internal_factory( + tmp_path, monkeypatch, mode +): + root = tmp_path / "source" + root.mkdir() + subprocess.run(["git", "init", str(root)], check=True, capture_output=True) + subprocess.run( + [ + "git", + "-C", + str(root), + "-c", + "user.name=CPU Control", + "-c", + "user.email=cpu-control@example.invalid", + "commit", + "--allow-empty", + "-m", + "CPU", + ], + check=True, + capture_output=True, + ) + model = tmp_path / "model" + model.mkdir() + (model / "config.json").write_text('{"dspark_block_size": 3}', encoding="utf-8") + (model / "model.safetensors.index.json").write_text( + '{"weight_map": {"mtp.0.weight": "not-a-real-model.safetensors"}}', encoding="utf-8" + ) + config = { + "skip_tokenizer_init": True, + "max_batch_size": 2, + "max_seq_len": 128, + "max_num_tokens": 16, + "disable_overlap_scheduler": False, + "cuda_graph_config": {"batch_sizes": [1, 2], "enable_padding": True}, + "speculative_config": { + "decoding_type": "DSpark", + "max_draft_len": 3, + "speculative_model": str(model), + }, + } + config_path = tmp_path / "runtime.yml" + config_path.write_text(yaml.safe_dump(config), encoding="utf-8") + inputs = tmp_path / "inputs.jsonl" + inputs.write_text( + "\n".join(json.dumps({"prompt_token_ids": [4, 5]}) for _ in range(2)), encoding="utf-8" + ) + cost = tmp_path / "structure-is-not-validated-at-argument-parser.json" + cost.write_text("{}", encoding="utf-8") + script = Path(dspark_reproduction.__file__).parents[2] / "tools/dspark_reproduction.py" + spec = importlib.util.spec_from_file_location("_actual_public_endpoint_cli", script) + cli = importlib.util.module_from_spec(spec) + spec.loader.exec_module(cli) + seen = [] + + def prepare_only(args): + llm_args, rows, contract, _ = dspark_reproduction._prepare(args, 1) + assert isinstance(llm_args, TorchLlmArgs) + assert llm_args.speculative_config.max_draft_len == 3 + assert llm_args.speculative_config.enable_confidence_scheduling == (mode == "dynamic") + # Regression: internal executor_cls MUST fail the genuine public LLM + # boundary, not be smuggled into public kwargs and discovered on a GPU. + public_kwargs = {**config, "model": str(model), "backend": "pytorch"} + LLM._validate_args_for_torch_backend(None, public_kwargs) + with pytest.raises(ValueError, match="executor_cls"): + LLM._validate_args_for_torch_backend(None, {**public_kwargs, "executor_cls": object}) + original = py_executor_creator.create_py_executor + + def capture_factory(**kwargs): + # functools.wraps exposes the inner implementation by default; + # bind the real callable, which owns its internal startup timer. + outer_signature = inspect.signature(original, follow_wrapped=False) + assert "_startup_timer" not in outer_signature.parameters + outer_signature.bind(**kwargs) + assert "executor_cls" not in kwargs + assert kwargs["llm_args"] is llm_args + assert isinstance( + kwargs["runtime_constructors"], py_executor_creator._RuntimeConstructors + ) + seen.append(kwargs) + inner = py_executor_creator._create_py_executor_impl + forwarded = [] + + def capture_model_construction(**arguments): + inspect.signature(inner, follow_wrapped=False).bind(**arguments) + timer = arguments["_startup_timer"] + assert isinstance(timer, py_executor_creator._StartupTimer) + assert timer.name == "executor_creation" and timer.started > 0 + assert { + key: value for key, value in arguments.items() if key != "_startup_timer" + } == { + "llm_args": llm_args, + "checkpoint_dir": kwargs["checkpoint_dir"], + "tokenizer": kwargs["tokenizer"], + "profiling_stage_data": None, + "resource_governor_queue": None, + "runtime_constructors": kwargs["runtime_constructors"], + } + forwarded.append(arguments) + return "CPU constructor boundary only" + + # Execute the genuine decorated wrapper and real host-only timer; + # intercept only the inner expensive model construction boundary. + monkeypatch.setattr( + py_executor_creator, "_create_py_executor_impl", capture_model_construction + ) + result = original(**kwargs) + assert len(forwarded) == 1 + return result + + monkeypatch.setattr(py_executor_creator, "create_py_executor", capture_factory) + assert ( + dspark_observation_runtime.create_observation_executor( + llm_args=llm_args, + checkpoint_dir=str(model), + tokenizer=None, + collect_calibration=mode == "collect", + ) + == "CPU constructor boundary only" + ) + assert len(rows) == 2 and contract["batch_size"] == 2 + + monkeypatch.setattr(dspark_reproduction, "run_native", prepare_only) + args = [ + "run", + "--mode", + mode, + "--model", + str(model), + "--config", + str(config_path), + "--inputs", + str(inputs), + "--source-root", + str(root), + "--physical-k", + "3", + "--batch-size", + "2", + "--output-tokens", + "8", + "--timeout", + "3600", + "--output", + str(tmp_path / "private-output"), + ] + if mode == "dynamic": + args += ["--cost-table", str(cost)] + cli.main(args) + assert len(seen) == 1 + # Keep the original setup (including MPI/device/loader handling), overriding + # only its existing internal factory seam. + assert dspark_reproduction.ObservationWorker.setup_engine is BaseWorker.setup_engine + + +@pytest.mark.cpu_only +def test_offline_head_evaluation_rejects_enabled_policy_before_factory(monkeypatch, tmp_path): + def forbidden(**kwargs): + raise AssertionError("must refuse before model construction") + + monkeypatch.setattr(py_executor_creator, "create_py_executor", forbidden) + from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig + + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"mtp.0.weight": "not-a-real-model.safetensors"}}', encoding="utf-8" + ) + spec = DSparkDecodingConfig( + max_draft_len=3, speculative_model=str(tmp_path), enable_confidence_scheduling=True + ) + with pytest.raises(ValueError, match="native full-K"): + dspark_observation_runtime.create_observation_executor( + llm_args=SimpleNamespace(speculative_config=spec), collect_calibration=True + ) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("resource_governor", [False, True]) +def test_profile_factory_keeps_explicit_offline_constructors( + monkeypatch: pytest.MonkeyPatch, resource_governor: bool +) -> None: + profile = object() + llm_args = SimpleNamespace(enable_resource_governor=resource_governor) + seen = [] + result = object() + + def capture_factory(**kwargs): + inspect.signature(py_executor_creator.create_py_executor, follow_wrapped=False).bind( + **kwargs + ) + seen.append(kwargs) + return result + + monkeypatch.setattr(dspark_request_profile, "create_py_executor", capture_factory) + if resource_governor: + with pytest.raises(ValueError, match="resource-governor"): + dspark_request_profile.create_profile_executor(llm_args=llm_args, profile=profile) + assert not seen + return + + assert ( + dspark_request_profile.create_profile_executor( + llm_args=llm_args, profile=profile, checkpoint_dir="model", tokenizer=None + ) + is result + ) + assert len(seen) == 1 + constructors = seen[0]["runtime_constructors"] + assert isinstance(constructors, py_executor_creator._RuntimeConstructors) + assert constructors.model_engine.func is dspark_request_profile.OwnedDSparkProfileEngine + assert constructors.model_engine.keywords == {"profile": profile} + assert constructors.executor is dspark_request_profile.OwnedDSparkProfileExecutor + assert constructors.kv_cache_creator is dspark_request_profile._ProfileKvCacheCreator + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("active_profile", [False, True], ids=["native", "measured"]) +@pytest.mark.parametrize("position_id_offsets", [None, object()], ids=["default", "offsets"]) +def test_profile_replay_forwards_parent_position_offsets( + monkeypatch, active_profile, position_id_offsets +): + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import DSparkCUDAGraphRunner + from tensorrt_llm._torch.pyexecutor.dspark_profile_engine import _DSparkProfileGraphs + + # A newly added parent keyword must be explicitly supported, with the same + # default, before the expensive model warmup reaches this override. + parent_parameters = inspect.signature(DSparkCUDAGraphRunner.replay).parameters + profile_parameters = inspect.signature(_DSparkProfileGraphs.replay).parameters + parent_keywords = { + name: parameter + for name, parameter in parent_parameters.items() + if parameter.kind is inspect.Parameter.KEYWORD_ONLY + } + assert "position_id_offsets" in parent_keywords + assert parent_keywords.keys() <= profile_parameters.keys() + for name, parameter in parent_keywords.items(): + assert profile_parameters[name].kind is inspect.Parameter.KEYWORD_ONLY + assert profile_parameters[name].default == parameter.default + assert profile_parameters["position_id_offsets"].default is None + + key, output, lease = (64, 3, False), object(), object() + attention, speculative = object(), object() + inputs = { + "attn_metadata": attention, + "spec_metadata": speculative, + "optional_metadata": object(), + } + calls, events = [], [] + + def replay(self, actual_key, current_inputs, *, position_id_offsets=None): + calls.append((self, actual_key, current_inputs, position_id_offsets)) + events.append("replay") + return output + + def stable(metadata, snapshot, *, context): + assert metadata is attention or metadata is speculative + assert context["rank"] == 0 + + monkeypatch.setattr(DSparkCUDAGraphRunner, "replay", replay) + monkeypatch.setattr( + _DSparkProfileGraphs, "_assert_graph_tensor_ptrs_stable", staticmethod(stable) + ) + runner = _DSparkProfileGraphs.__new__(_DSparkProfileGraphs) + runner.active_profile_cell = ( + SimpleNamespace(graph_batch_size=64, verifier_budget=0) if active_profile else None + ) + runner.config = SimpleNamespace(mapping=SimpleNamespace(rank=0)) + runner._profile = SimpleNamespace(max_samples=4, physical_k=3) + runner._profile_records = [] + runner._profile_ptrs = {key: ({}, {})} + runner.graph_metadata = {key: inputs} + runner.active_request_ids = ("request",) + runner._profile_pool = SimpleNamespace( + begin=lambda: events.append("begin") or lease, + end=lambda actual: events.append(("end", actual is lease)), + ) + assert runner.replay(key, inputs, position_id_offsets=position_id_offsets) is output + assert len(calls) == 1 + assert calls[0][0] is runner and calls[0][1] is key + assert calls[0][2] is inputs and calls[0][3] is position_id_offsets + assert events == (["begin", "replay", ("end", True)] if active_profile else ["replay"]) + assert len(runner._profile_records) == int(active_profile) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("enabled", [False, True], ids=["native-collect", "dspark"]) +@pytest.mark.parametrize("position_id_offsets", [None, object()], ids=["default", "offsets"]) +def test_observation_replay_forwards_actual_parent_keyword_contract( + tmp_path, monkeypatch, enabled, position_id_offsets +): + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( + CUDAGraphRunner, + DSparkCUDAGraphRunner, + KeyType, + ) + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner + from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine + from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig + + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"mtp.0.weight": "not-a-real-model.safetensors"}}', encoding="utf-8" + ) + spec = DSparkDecodingConfig( + max_draft_len=3, + speculative_model=str(tmp_path), + enable_confidence_scheduling=enabled, + ) + captured = [] + + def capture_factory(**kwargs): + captured.append(kwargs["runtime_constructors"]) + return "CPU factory boundary only" + + monkeypatch.setattr(py_executor_creator, "create_py_executor", capture_factory) + assert ( + dspark_observation_runtime.observation_factory(collect_calibration=not enabled)( + llm_args=SimpleNamespace(speculative_config=spec), tokenizer=None + ) + == "CPU factory boundary only" + ) + engine_class = captured[0].model_engine + monkeypatch.setattr( + PyTorchModelEngine, + "_initialize_decoder_runner", + lambda self, runner_cls, runner_config=None, **kwargs: runner_cls, + ) + engine = engine_class.__new__(engine_class) + decoder_class = engine._initialize_decoder_runner(DecoderRunner) + base = DSparkCUDAGraphRunner if enabled else CUDAGraphRunner + parent_parameters = inspect.signature(base.replay).parameters + monkeypatch.setattr(base, "__init__", lambda self, config: None) + decoder = decoder_class.__new__(decoder_class) + graph = decoder._create_cuda_graph_runner(SimpleNamespace()) + assert type(graph).__mro__[1] is base + parameters = inspect.signature(type(graph).replay).parameters + keywords = { + name: parameter + for name, parameter in parent_parameters.items() + if parameter.kind is inspect.Parameter.KEYWORD_ONLY + } + assert "position_id_offsets" in keywords + assert keywords.keys() <= parameters.keys() + for name, parameter in keywords.items(): + assert parameters[name].kind is inspect.Parameter.KEYWORD_ONLY + assert parameters[name].default == parameter.default + + key = KeyType(64, 3, False, ragged_verify_bucket=64 if enabled else None) + inputs = {"attn_metadata": object(), "spec_metadata": object(), "optional": object()} + output, calls = object(), [] + + def replay(self, actual_key, actual_inputs, *, position_id_offsets=None): + calls.append((self, actual_key, actual_inputs, position_id_offsets)) + return output + + monkeypatch.setattr(base, "replay", replay) + assert graph.replay(key, inputs, position_id_offsets=position_id_offsets) is output + assert len(calls) == 1 + assert calls[0][0] is graph and calls[0][1] is key + assert calls[0][2] is inputs and calls[0][3] is position_id_offsets + assert graph.observed_routes == { + "native_replays": int(not enabled), + "compact_replays": int(enabled), + } + assert len(graph.observed_keys) == 1 + with pytest.raises(TypeError): + graph.replay(key, inputs, unsupported_next_keyword=object()) diff --git a/tests/unittest/tools/test_dspark_calibration.py b/tests/unittest/tools/test_dspark_calibration.py new file mode 100644 index 000000000000..1ab167646704 --- /dev/null +++ b/tests/unittest/tools/test_dspark_calibration.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Actual CPU STS maths, input contracts and standalone CLI controls.""" + +import hashlib +import importlib.util +import json +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +import torch + +REPO = Path(__file__).resolve().parents[3] +SCRIPT = REPO / "tensorrt_llm/tools/dspark_calibration.py" +SPEC = importlib.util.spec_from_file_location("standalone_dspark_calibration", SCRIPT) +calibration = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(calibration) + + +class CalibrationCpuTests(unittest.TestCase): + def test_known_cumulative_prefix_target_and_grid_order(self) -> None: + logits = torch.zeros((4, 2), dtype=torch.float32) + labels = torch.tensor([[1, 1], [1, 0], [0, 0], [0, 0]], dtype=torch.float32) + grid = calibration.default_temperature_grid() + result = calibration.fit_sts_temperatures(logits=logits, prefix_mask=labels, grid=grid) + self.assertEqual(grid.numel(), 41) + self.assertEqual(result["temperatures"], [float(grid[0]), float(grid[0])]) + self.assertEqual(result["ece_before"], [0.0, 0.0]) + self.assertEqual(result["ece_after"], [0.0, 0.0]) + + def test_weighted_ece_not_unweighted_bin_mean(self) -> None: + probs = torch.tensor([0.1, 0.1, 0.1, 0.9]) + targets = torch.tensor([0.0, 0.0, 0.0, 0.0]) + expected = float((probs.to(torch.float64).sum() / 4).item()) + self.assertAlmostEqual( + calibration.expected_calibration_error(probs=probs, targets=targets, num_bins=10), + expected, + ) + + def test_physical_k_is_not_fixed_to_pilot_three(self) -> None: + for width in (1, 2, 3, 4, 5, 8): + with self.subTest(width=width): + logits = torch.zeros((4, width)) + labels = torch.zeros((4, width)) + result = calibration.fit_sts_temperatures( + logits=logits, prefix_mask=labels, grid=calibration.default_temperature_grid() + ) + self.assertEqual(len(result["temperatures"]), width) + self.assertTrue(all(value > 0 for value in result["temperatures"])) + + def test_malformed_rows_and_censored_prefix_order_refused(self) -> None: + grid = calibration.default_temperature_grid() + rows = [ + (torch.empty((0, 3)), torch.empty((0, 3))), + (torch.zeros((2, 3)), torch.zeros((2, 2))), + (torch.zeros((2, 3), dtype=torch.int32), torch.zeros((2, 3))), + (torch.tensor([[float("nan")]]), torch.zeros((1, 1))), + (torch.zeros((1, 1)), torch.tensor([[float("inf")]])), + (torch.zeros((1, 1)), torch.tensor([[0.5]])), + (torch.zeros((1, 3)), torch.tensor([[1, 0, 1]])), + ] + for logits, labels in rows: + with self.subTest(shape=list(logits.shape)), self.assertRaises(ValueError): + calibration.fit_sts_temperatures(logits=logits, prefix_mask=labels, grid=grid) + for bad in (torch.tensor([]), torch.tensor([0.0]), torch.tensor([float("inf")])): + with self.subTest(grid=bad), self.assertRaises(ValueError): + calibration.fit_sts_temperatures( + logits=torch.zeros((2, 1)), prefix_mask=torch.zeros((2, 1)), grid=bad + ) + with self.assertRaises(ValueError): + calibration.fit_sts_temperatures( + logits=torch.zeros((2, 1)), prefix_mask=torch.zeros((2, 1)), grid=grid, num_bins=0 + ) + + def test_safe_shard_pairing_width_and_duplicate_contracts(self) -> None: + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "paired.pt" + logits = torch.zeros((4, 3)) + mask = torch.zeros((4, 3)) + torch.save( + {"meta": {"pairing": "draft_seq_ring"}, "logits": logits, "prefix_mask": mask}, path + ) + actual_logits, actual_mask, hashes = calibration.load_shards([path], 3) + self.assertTrue(torch.equal(logits, actual_logits)) + self.assertTrue(torch.equal(mask, actual_mask)) + self.assertEqual(hashes[str(path)], hashlib.sha256(path.read_bytes()).hexdigest()) + with self.assertRaisesRegex(ValueError, "Duplicate"): + calibration.load_shards([path, path], 3) + with self.assertRaisesRegex(ValueError, "width"): + calibration.load_shards([path], 5) + torch.save( + { + "meta": {"pairing": "unpaired_live_buffer"}, + "logits": logits, + "prefix_mask": mask, + }, + path, + ) + with self.assertRaisesRegex(ValueError, "same stamped draft"): + calibration.load_shards([path], 3) + torch.save({"logits": logits, "prefix_mask": mask}, path) + with self.assertRaisesRegex(ValueError, "pairing metadata"): + calibration.load_shards([path], 3) + + def test_real_cli_from_unrelated_cwd_and_exclusive_output(self) -> None: + with tempfile.TemporaryDirectory() as directory: + base = Path(directory) + shards, output = base / "paired.pt", base / "sts.json" + logits = torch.zeros((4, 3)) + mask = torch.tensor([[1, 1, 1], [1, 1, 0], [1, 0, 0], [0, 0, 0]], dtype=torch.float32) + torch.save( + {"meta": {"pairing": "draft_seq_ring"}, "logits": logits, "prefix_mask": mask}, + shards, + ) + argv = [ + sys.executable, + "-I", + "-B", + str(SCRIPT), + "--shards", + str(shards), + "--physical-k", + "3", + "--output", + str(output), + "--minimum-samples", + "4", + ] + result = subprocess.run(argv, cwd=directory, capture_output=True, text=True, timeout=60) + self.assertEqual(result.returncode, 0, result.stderr) + payload = json.loads(output.read_bytes()) + self.assertEqual(payload["sts_temperatures"], payload["temperatures"]) + self.assertEqual(payload["num_samples"], 4) + self.assertEqual(payload["physical_k"], 3) + self.assertFalse(payload["task_accuracy_or_economic_admission_qualified"]) + self.assertFalse(payload["artifact_freshness_validated"]) + before = output.read_bytes() + again = subprocess.run(argv, cwd=directory, capture_output=True, text=True, timeout=60) + self.assertNotEqual(again.returncode, 0) + self.assertIn("FileExistsError", again.stderr) + self.assertEqual(output.read_bytes(), before) + + def test_no_native_trtllm_import_and_output_does_not_fabricate_admission(self) -> None: + text = SCRIPT.read_text() + self.assertNotIn("import tensorrt_llm", text) + self.assertNotIn("torch.cuda", text) + self.assertIn("weights_only=True", text) + self.assertNotIn('"admitted": True', text) + self.assertNotIn("/scratch/", text) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/unittest/tools/test_dspark_cost_export.py b/tests/unittest/tools/test_dspark_cost_export.py new file mode 100644 index 000000000000..87066313a6ad --- /dev/null +++ b/tests/unittest/tools/test_dspark_cost_export.py @@ -0,0 +1,409 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU artifact controls, not measured GPU costs or performance evidence.""" + +import copy +import hashlib +import importlib.util +import json +import statistics +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[3] +SOURCE = ROOT / "tensorrt_llm/tools/dspark_cost_export.py" +SPEC = importlib.util.spec_from_file_location("public_dspark_cost_export_cpu", SOURCE) +MODULE = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = MODULE +SPEC.loader.exec_module(MODULE) + + +def _digest(value: object) -> str: + return hashlib.sha256( + json.dumps(value, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + + +def _costs(k: int = 3, graphs: tuple[int, ...] = (16, 64), world: int = 8) -> dict: + """Typed synthetic fixtures only; these numbers are not benchmark prices.""" + fingerprint = { + "global_graph_batch_sizes": [g * world for g in graphs], + "gpu": "CPU_SCHEMA_FIXTURE_NOT_GPU", + "gpu_count": world, + "gpu_snapshot_sha256": "a" * 64, + "max_draft_len": k, + "rank_local_graph_batch_sizes": list(graphs), + "runtime_snapshot": "CPU_FIXTURE_NOT_NATIVE", + "source_diff_sha256": "b" * 64, + "source_head": "ORIGINAL_RECORDED_CPU_FIXTURE_SOURCE", + "topology": "CPU_FIXTURE_NOT_FABRIC", + } + tables, measurements = {}, [] + for g in graphs: + budgets = [0] if k == 1 else [0, g, g * 2] + values = [float(index + 1) for index in range(len(budgets))] + tables[str(g)] = {"token_counts": budgets, "step_time_ms": values} + measurements.extend( + { + "rank_local_graph_batch_size": g, + "rank_local_verifier_budget": v, + "source_result_sha256": "c" * 64, + "step_time_ms": value, + } + for v, value in zip(budgets, values) + ) + return { + "schema_version": 2, + "cost_tables": tables, + "engine_fingerprint": fingerprint, + "engine_fingerprint_sha256": _digest(fingerprint), + "minimum_predicted_gain": 0.01, + "measurements": measurements, + } + + +def _check(payload: dict, k: int = 3, graphs: tuple[int, ...] = (16, 64)) -> dict: + return MODULE.check_cost_bytes( + json.dumps(payload).encode(), physical_k=k, graph_batch_sizes=list(graphs) + ) + + +class PublicCostExportTests(unittest.TestCase): + def _whole_step_summary(self, k: int = 3) -> dict: + """Exercise the real reducer on explicitly synthetic CPU raw records.""" + utilities = MODULE._source_module( + "cost_bridge_test_reducer", ROOT / "tensorrt_llm/tools/verify_costs.py" + ) + costs = _costs(k=k, graphs=(16,), world=2) + declared = [] + shapes = [] + for v in costs["cost_tables"]["16"]["token_counts"]: + shape = { + "workload": "dspark_chain", + "physical_k": k, + "graph_batch_size": 16, + "verifier_budget": v, + "attention_layout": "uniform" if v == 0 else "packed_chain_rows", + "captured_head": True, + } + shapes.append(shape) + declared.append(utilities.canonical_sha256(shape)) + identity = { + "cpu_fixture_not_gpu": True, + "generations_per_shape": 3, + "declared_shape_ids": declared, + "recorded_engine_fingerprint": costs["engine_fingerprint"], + "profile": { + "physical_k": k, + "record_confidence_head": True, + "measurement_scope": "whole-step", + "cells": [ + {"graph_batch_size": 16, "verifier_budget": shape["verifier_budget"]} + for shape in shapes + ], + }, + } + raw = [] + for cell_index, (shape, declared_id) in enumerate(zip(shapes, declared)): + for generation in range(3): + for rank in range(2): + raw.append( + { + "schema": "verify-cost-raw-v1", + "rank": rank, + "world_size": 2, + "generation": generation, + "generation_closed": True, + "run_complete": True, + "model_shutdown_completed": True, + "identity": identity, + "timing_scope": MODULE.WHOLE_STEP_SCOPE, + "declared_shape_id": declared_id, + "execution_shape": {**shape, "observed_graph_key": [16, k]}, + "samples": [ + { + "sample_index": cell_index * 42 + generation * 14 + i, + "executor_iteration": cell_index * 42 + + generation * 14 + + i + + 100, + "elapsed_ms": float(rank + i + cell_index + 1), + "eligible": True, + } + for i in range(14) + ], + } + ) + return utilities.reduce_measurements(raw) + + def test_complete_loop_reducer_to_actual_consumer_preserves_values(self) -> None: + for k in (1, 2, 3, 5, 8): + summary = self._whole_step_summary(k) + before = copy.deepcopy(summary) + data = json.dumps(summary).encode() + output, report = MODULE.runtime_cost_bytes(data, physical_k=k, graph_batch_sizes=[16]) + payload = json.loads(output) + self.assertEqual(summary, before) + self.assertEqual( + payload["engine_fingerprint"], summary["identity"]["recorded_engine_fingerprint"] + ) + self.assertTrue(report["direct_whole_step_measurements_exported"]) + self.assertFalse(report["performance_or_economic_admission_qualified"]) + self.assertNotIn("iteration_drain_guard", payload) + expected = { + cell["execution_shape"]["verifier_budget"]: statistics.median( + cell["retained_rank_max_samples_ms"] + ) + for cell in summary["cells"] + } + for row in payload["measurements"]: + self.assertEqual(row["step_time_ms"], expected[row["rank_local_verifier_budget"]]) + self.assertEqual(row["source_result_sha256"], hashlib.sha256(data).hexdigest()) + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "costs.json" + path.write_bytes(output) + consumer = MODULE._source_module( + "bridge_test_actual_consumer", + ROOT / "tensorrt_llm/_torch/speculative/dspark_planner.py", + ) + table, _ = consumer.load_runtime_sps_cost_table( + path, max_draft_len=k, graph_batch_sizes=[16], validate_freshness=False + ) + self.assertEqual(table.step_time(0, 16), expected[0]) + + def test_complete_loop_export_refuses_missing_or_repriced_cells(self) -> None: + for kind in ("missing", "median", "fingerprint", "K", "layout", "duplicate", "floor"): + summary = self._whole_step_summary() + if kind == "missing": + summary["cells"].pop() + elif kind == "median": + summary["cells"][0]["rank_max_median_ms"] += 1.0 + elif kind == "fingerprint": + del summary["identity"]["recorded_engine_fingerprint"] + elif kind == "K": + summary["cells"][0]["execution_shape"]["physical_k"] = 5 + elif kind == "layout": + summary["cells"][0]["execution_shape"]["attention_layout"] = "other" + elif kind == "duplicate": + summary["cells"].append(copy.deepcopy(summary["cells"][0])) + else: + summary["sample_floors"]["pooled_per_shape"] = 1 + with self.subTest(kind=kind), self.assertRaises((ValueError, KeyError)): + MODULE.runtime_cost_bytes( + json.dumps(summary).encode(), physical_k=3, graph_batch_sizes=[16] + ) + + def test_existing_costs_use_actual_consumer_without_freshness_gate(self) -> None: + for k in (1, 2, 3, 4, 5, 8): + for world in (1, 2, 8): + with self.subTest(k=k, world=world): + payload = _costs(k=k, graphs=(8, 192, 256), world=world) + before = copy.deepcopy(payload) + report = _check(payload, k=k, graphs=(8, 192, 256)) + subset_report = _check(payload, k=k, graphs=(8,)) + self.assertTrue(subset_report["structural_validation_passed"]) + self.assertEqual(payload, before) + self.assertTrue(report["structural_validation_passed"]) + self.assertFalse(report["destination_freshness_checked"]) + self.assertFalse(report["new_measurements_created"]) + self.assertFalse(report["performance_or_economic_admission_qualified"]) + self.assertEqual( + report["recorded_source_head"], + payload["engine_fingerprint"]["source_head"], + ) + + def test_optional_measured_drain_values_retained_without_invention(self) -> None: + payload = _costs() + payload["iteration_drain_guard"] = { + "loss_multiplier": 1.5, + "mean_output_tokens_per_request_iteration": 3.25, + "minimum_group_value_ms": 0.2, + "source_result_sha256": "d" * 64, + "tail_graph_batch_size": 16, + } + before = copy.deepcopy(payload) + _check(payload) + self.assertEqual(payload, before) + self.assertNotIn("iteration_drain_guard", _costs()) + + def test_actual_runtime_candidate_bounds_refuse_before_copy(self) -> None: + payload = _costs(graphs=(16,)) + budgets = [0, 16, 17, 18, 19, 20] + payload["cost_tables"]["16"] = { + "token_counts": budgets, + "step_time_ms": [1.0] * len(budgets), + } + payload["measurements"] = [ + { + "rank_local_graph_batch_size": 16, + "rank_local_verifier_budget": v, + "source_result_sha256": "c" * 64, + "step_time_ms": 1.0, + } + for v in budgets + ] + with self.assertRaisesRegex(ValueError, "compact V cells"): + _check(payload, graphs=(16,)) + + def test_actual_consumer_rejects_invalid_schema_integrity_or_cells(self) -> None: + for kind in ( + "physical_k", + "graph_coverage", + "fingerprint", + "measurement", + "native", + "negative", + "extra", + "bool", + "unknown_scope", + ): + payload = _costs() + k, graphs = 3, (16, 64) + if kind == "physical_k": + k = 5 + elif kind == "graph_coverage": + graphs = (16, 32) + elif kind == "fingerprint": + payload["engine_fingerprint_sha256"] = "f" * 64 + elif kind == "measurement": + payload["measurements"][0]["step_time_ms"] = 99.0 + elif kind == "native": + payload["cost_tables"]["16"]["token_counts"][0] = 1 + elif kind == "negative": + payload["cost_tables"]["16"]["step_time_ms"][0] = -1.0 + elif kind == "extra": + payload["measurements"].append(copy.deepcopy(payload["measurements"][0])) + elif kind == "bool": + payload["cost_tables"]["16"]["token_counts"][1] = True + else: + payload["timing_scope"] = "partial-scope-cannot-be-hidden" + with self.subTest(kind=kind), self.assertRaises((TypeError, ValueError)): + _check(payload, k=k, graphs=graphs) + + def test_partial_summary_refuses_even_if_eligibility_is_toggled(self) -> None: + for schema in ("verify-cost-raw-v1", "verify-cost-measured-summary-v1"): + for eligible in (False, True): + payload = { + "schema": schema, + "timing_scope": "original-graph-replay-entry-including-static-copies", + "identity": {"identity_scope": "DIAGNOSTIC_PARTIAL_LOADED_RUNTIME"}, + "production_sps_table_eligible": eligible, + } + with self.subTest(schema=schema, eligible=eligible): + with self.assertRaisesRegex(ValueError, "whole-step producer"): + _check(payload) + + def test_ambiguous_or_nonfinite_json_refuses_before_export(self) -> None: + for data in (b'{"schema_version":2,"schema_version":2}', b'{"value":NaN}', b"[]"): + with self.subTest(data=data), self.assertRaises((ValueError, TypeError)): + MODULE.check_cost_bytes(data, physical_k=3, graph_batch_sizes=[16, 64]) + + def test_real_standalone_cli_preserves_bytes_and_refuses_overwrite(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source, output = root / "costs.json", root / "new-costs.json" + original = json.dumps(_costs(), indent=3).encode() + b"\n\n" + source.write_bytes(original) + command = [ + sys.executable, + "-I", + "-B", + str(SOURCE), + "--input", + str(source), + "--physical-k", + "3", + "--graph-batch-sizes", + "16", + "64", + "--output", + str(output), + ] + result = subprocess.run(command, cwd=root, capture_output=True, timeout=30) + self.assertEqual(result.returncode, 0, result.stderr.decode()) + self.assertEqual(source.read_bytes(), original) + self.assertEqual(output.read_bytes(), original) + report = json.loads(result.stdout) + self.assertEqual(report["source_artifact_sha256"], hashlib.sha256(original).hexdigest()) + again = subprocess.run(command, cwd=root, capture_output=True, timeout=30) + self.assertNotEqual(again.returncode, 0) + self.assertEqual(output.read_bytes(), original) + + def test_real_complete_loop_cli_writes_only_original_direct_cells(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source, output = root / "CPU-summary.json", root / "CPU-costs.json" + summary = self._whole_step_summary(k=5) + source.write_text(json.dumps(summary), encoding="utf-8") + result = subprocess.run( + [ + sys.executable, + "-I", + "-B", + str(SOURCE), + "--input", + str(source), + "--physical-k", + "5", + "--graph-batch-sizes", + "16", + "--output", + str(output), + ], + cwd=root, + capture_output=True, + timeout=30, + ) + self.assertEqual(result.returncode, 0, result.stderr.decode()) + payload = json.loads(output.read_bytes()) + self.assertEqual( + payload["engine_fingerprint"], summary["identity"]["recorded_engine_fingerprint"] + ) + self.assertEqual(payload["minimum_predicted_gain"], 0.01) + self.assertFalse( + json.loads(result.stdout)["performance_or_economic_admission_qualified"] + ) + + def test_real_partial_export_cli_creates_no_cost_artifact(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source, output = root / "partial.json", root / "must-not-exist.json" + source.write_text(json.dumps({"schema": "verify-cost-measured-summary-v1"})) + result = subprocess.run( + [ + sys.executable, + "-I", + "-B", + str(SOURCE), + "--input", + str(source), + "--physical-k", + "3", + "--graph-batch-sizes", + "16", + "64", + "--output", + str(output), + ], + cwd=root, + capture_output=True, + timeout=30, + ) + self.assertNotEqual(result.returncode, 0) + self.assertIn(b"whole-step producer", result.stderr) + self.assertFalse(output.exists()) + + def test_cpu_dependency_loading_never_initializes_native_package(self) -> None: + _check(_costs()) + self.assertNotIn("tensorrt_llm", sys.modules) + self.assertNotIn("torch", sys.modules) + source = SOURCE.read_text(encoding="utf-8") + for token in ("/scratch/", "aws-pdx", "K3", "admitted=True"): + self.assertNotIn(token, source) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/unittest/tools/test_dspark_observations.py b/tests/unittest/tools/test_dspark_observations.py new file mode 100644 index 000000000000..9fbc2367f47b --- /dev/null +++ b/tests/unittest/tools/test_dspark_observations.py @@ -0,0 +1,445 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Source/CPU controls only; finite parents do not qualify native overlap/CUDA.""" + +import copy +import importlib.util +import json +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace + +import torch + +ROOT = Path(__file__).resolve().parents[3] +PUBLIC = ROOT / "tensorrt_llm/tools" + + +def load(name, path, source=None): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + if source is None: + spec.loader.exec_module(module) + else: + exec(compile(source, str(path), "exec"), module.__dict__) + return module + + +observations = load("tensorrt_llm.tools.dspark_observations", PUBLIC / "dspark_observations.py") +RUNTIME = ROOT / "tensorrt_llm/_torch/pyexecutor/dspark_observation_runtime.py" +runtime = load("_public_observation_runtime", RUNTIME) + + +def request(identifier, slot): + return SimpleNamespace( + py_request_id=identifier, + py_seq_slot=slot, + py_decoding_iter=0, + is_dummy=False, + py_draft_tokens=[7, 8, 9], + ) + + +class FiniteSampler: + """Deferred copies covered by the finite ORIGINAL sampler event, not a GPU event.""" + + def __init__(self): + self.copies = [] + self.waits = 0 + + def _copy_to_host(self, tensor): + host = torch.zeros_like(tensor, device="cpu") + self.copies.append((host, tensor.clone())) + return host + + def event(self): + copies, self.copies = self.copies, [] + + def synchronize(): + self.waits += 1 + for host, snapshot in copies: + host.copy_(snapshot) + + return SimpleNamespace(synchronize=synchronize) + + +class FiniteWorker: + def __init__(self): + self.sequence = 0 + self.confidence = torch.zeros((4, 3)) + self.stamps = torch.zeros(4, dtype=torch.int32) + + def bump_draft_seq(self): + old = self.sequence + self.sequence += 1 + return old + + def staged_confidence_buffer(self): + return self.confidence + + def confidence_stamp_buffer(self): + return self.stamps + + def confidence_row_for(self, identifier): + return {41: 0, 42: 3}[identifier] + + +class FiniteOriginal: + """Only the parent's finite call/order/buffer contract is modeled here.""" + + def __init__(self): + self.is_warmup = False + self.iter_counter = 10 + self.worker = FiniteWorker() + self.model_engine = SimpleNamespace( + _collect_calibration=True, + model=SimpleNamespace(spec_worker=self.worker), + spec_config=SimpleNamespace(max_draft_len=3), + ) + self.sampler = FiniteSampler() + self.tokens = torch.zeros((3, 3), dtype=torch.int32) + + def _forward_step(self, batch, new_tensors_device=None): + for req in batch.generation_requests: + row = self.worker.confidence_row_for(req.py_request_id) + self.worker.confidence[row] = torch.tensor( + [ + self.iter_counter + row, + self.iter_counter + row + 0.1, + self.iter_counter + row + 0.2, + ] + ) + self.worker.stamps[row] = self.worker.sequence + self.tokens[req.py_seq_slot] = torch.tensor( + [self.iter_counter * 10 + i for i in range(3)] + ) + return {"next_draft_tokens": self.tokens} + + def _sample_async(self, batch, output): + host = SimpleNamespace( + next_draft_tokens=self.sampler._copy_to_host(self.tokens), + new_tokens_lens=self.sampler._copy_to_host(torch.tensor([2, 1, 3])), + verify_lens=None, + ) + return SimpleNamespace( + host=host, + device=SimpleNamespace(next_draft_tokens=self.tokens), + sampler_event=self.sampler.event(), + requests=list(batch.generation_requests), + runtime_draft_len=3, + draft_lens=[3] * len(batch.generation_requests), + ) + + def _update_requests(self, state, resource_manager=None): + state.sampler_event.synchronize() + for req in state.requests: + if getattr(req, "retired", False): + continue + req.py_decoding_iter += 1 + req.py_draft_tokens = state.host.next_draft_tokens[req.py_seq_slot].tolist() + + +def complete_overlap(mixin): + class Executor(mixin, FiniteOriginal): + pass + + executor = Executor() + reqs = [request(41, 2), request(42, 0)] + batch = SimpleNamespace(generation_requests=reqs) + output = executor._forward_step(batch) + source = executor._sample_async(batch, output) + executor.iter_counter = 11 + reordered = SimpleNamespace(generation_requests=list(reversed(reqs))) + output = executor._forward_step(reordered, source.device) + executor._update_requests(source) + target = executor._sample_async(reordered, output) + executor._update_requests(target) + return executor + + +def run_record(mode): + inputs = [{"input_id": 0, "prompt_token_ids": [4, 5]}] + return { + "schema": "dspark-public-run-v1", + "mode": mode, + "run_complete": True, + "model_shutdown_completed": True, + "calibration_collection": False, + "comparison_contract": { + "physical_k": 3, + "batch_size": 1, + "output_tokens": 2, + "same_cpu_fixture": True, + }, + "elapsed_seconds": 2 if mode == "static" else 1.5, + "resolved_confidence_enabled": mode == "dynamic", + "runtime_confidence_enabled": mode == "dynamic", + "routing": {"native_replays": 2, "compact_replays": 1 if mode == "dynamic" else 0}, + "inputs": inputs, + "responses": [{**inputs[0], "output_token_ids": [6, 7], "finish_reason": "length"}], + } + + +class Tests(unittest.TestCase): + def test_whole_overlap_reordered_slots_and_delayed_original_event(self): + executor = complete_overlap(runtime.ObservationExecutorMixin) + payload = executor._native_pairs.payload() + self.assertEqual(executor.sampler.waits, 2) + self.assertEqual( + payload["logits"].tolist()[0], [13.0, 13.100000381469727, 13.199999809265137] + ) + self.assertEqual( + payload["logits"].tolist()[1], [10.0, 10.100000381469727, 10.199999809265137] + ) + self.assertEqual(payload["prefix_mask"].tolist(), [[1, 0, 0], [1, 1, 0]]) + self.assertTrue(all(join["pair_key"][2:] == [10, 11] for join in payload["meta"]["joins"])) + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "pairs.pt" + torch.save(payload, path) + loaded = torch.load(path, map_location="cpu", weights_only=True) + self.assertEqual(loaded["meta"]["pairing"], "draft_seq_ring") + + def test_whole_old_fixed_restored_persistent_input_overwrite(self): + source = RUNTIME.read_text(encoding="utf-8") + before = "input_host = self.sampler._copy_to_host(new_tensors_device.next_draft_tokens)" + old = source.replace(before, "input_host = new_tensors_device.next_draft_tokens", 1) + for index, variant in enumerate((old, source, old)): + loaded = load(f"_source_boundary_{index}", RUNTIME, variant) + if index == 1: + self.assertEqual( + complete_overlap(loaded.ObservationExecutorMixin)._native_pairs.stats[ + "recorded" + ], + 2, + ) + else: + with self.assertRaisesRegex(ValueError, "verification input differs"): + complete_overlap(loaded.ObservationExecutorMixin) + + def test_wrong_step_incarnation_stale_confidence_and_censored_negatives(self): + executor = complete_overlap(runtime.ObservationExecutorMixin) + source = next(iter(executor._observation_states.values())) + req = source.requests[41][0] + arguments = dict( + source=source, + request=req, + target_iteration=12, + input_proposal=source.state.host.next_draft_tokens[2], + accepted=2, + verified=3, + target_slot=2, + ) + for fields, message in ( + ({"target_iteration": 10}, "after"), + ({"request": request(41, 2)}, "incarnation"), + ({"verified": 2}, "censors"), + ({"accepted": 4}, "outside"), + ): + with self.assertRaisesRegex(ValueError, message): + observations.NativePairs(3).record(**{**arguments, **fields}) + source.stamp_host[0] += 1 + with self.assertRaisesRegex(ValueError, "producing draft"): + observations.NativePairs(3).record(**arguments) + + def test_retired_target_is_not_labeled(self): + class Executor(runtime.ObservationExecutorMixin, FiniteOriginal): + pass + + executor = Executor() + req = request(41, 2) + batch = SimpleNamespace(generation_requests=[req]) + state = executor._sample_async(batch, executor._forward_step(batch)) + executor._update_requests(state) + executor.iter_counter += 1 + output = executor._forward_step(batch, state.device) + target = executor._sample_async(batch, output) + req.retired = True + executor._update_requests(target) + self.assertEqual(executor._native_pairs.stats["recorded"], 0) + self.assertEqual(executor._native_pairs.stats["retired_or_not_advanced"], 1) + + def test_same_state_cannot_update_twice(self): + executor = complete_overlap(runtime.ObservationExecutorMixin) + state = list(executor._observation_states.values())[-1].state + with self.assertRaisesRegex(RuntimeError, "consumed"): + executor._update_requests(state) + + def test_configurable_physical_k(self): + for k in (1, 2, 3, 4, 5, 8): + req = request(41, 0) + tokens = torch.arange(k).view(1, -1) + state = SimpleNamespace(host=SimpleNamespace(next_draft_tokens=tokens)) + snapshot = observations.ProposalSnapshot( + 1, state, torch.zeros((1, k)), {41: (req, 0, 0)}, torch.tensor([8]), 8, True + ) + pairs = observations.NativePairs(k) + pairs.record( + source=snapshot, + request=req, + target_iteration=2, + input_proposal=tokens[0], + accepted=k - 1, + verified=k, + target_slot=0, + ) + self.assertEqual(pairs.payload()["prefix_mask"].tolist(), [[1] * (k - 1) + [0]]) + + def test_prefill_placeholder_is_not_labeled_as_scored_draft(self): + class Executor(runtime.ObservationExecutorMixin, FiniteOriginal): + pass + + executor = Executor() + req = request(41, 2) + empty = SimpleNamespace(generation_requests=[]) + source = executor._sample_async(empty, executor._forward_step(empty)) + executor.iter_counter += 1 + batch = SimpleNamespace(generation_requests=[req]) + output = executor._forward_step(batch, source.device) + executor._update_requests(source) + target = executor._sample_async(batch, output) + executor._update_requests(target) + self.assertEqual(executor._native_pairs.stats["recorded"], 0) + self.assertEqual(executor._native_pairs.stats["prefill_placeholder_no_scored_proposal"], 1) + + def test_literal_collect_payload_to_approved_public_fitter(self): + fitter = PUBLIC / "dspark_calibration.py" + self.assertTrue(fitter.is_file(), "Assemble the approved public fitter with this endpoint") + with tempfile.TemporaryDirectory() as directory: + cwd = Path(directory) + shard = cwd / "cpu-source-contract-pairs.pt" + torch.save( + complete_overlap(runtime.ObservationExecutorMixin)._native_pairs.payload(), shard + ) + output = cwd / "cpu-source-contract-temperatures.json" + result = subprocess.run( + [ + sys.executable, + "-I", + "-B", + str(fitter), + "--shards", + str(shard), + "--physical-k", + "3", + "--minimum-samples", + "2", + "--output", + str(output), + ], + cwd=cwd, + capture_output=True, + text=True, + timeout=60, + ) + self.assertEqual(result.returncode, 0, result.stderr) + payload = json.loads(output.read_bytes()) + self.assertEqual(payload["physical_k"], 3) + self.assertEqual(payload["num_samples"], 2) + self.assertFalse(payload["task_accuracy_or_economic_admission_qualified"]) + + def test_complete_public_run_parser_has_configurable_k_bs_and_no_internal_kwargs(self): + cli = load("_public_cli_parser_control", PUBLIC / "dspark_reproduction.py") + name = "tensorrt_llm._torch.speculative.dspark_reproduction" + original = sys.modules.get(name) + seen = [] + sys.modules[name] = SimpleNamespace(run_native=lambda args: seen.append(args)) + try: + for mode in ("static", "dynamic", "collect"): + cli.main( + [ + "run", + "--mode", + mode, + "--model", + "model", + "--config", + "config.yml", + "--inputs", + "input.jsonl", + "--source-root", + "checkout", + "--physical-k", + "4", + "--batch-size", + "128", + "--output-tokens", + "64", + "--timeout", + "25200", + "--output", + "private-output", + ] + ) + self.assertEqual([args.mode for args in seen], ["static", "dynamic", "collect"]) + self.assertTrue(all(args.physical_k == 4 and args.batch_size == 128 for args in seen)) + self.assertTrue(all("executor_cls" not in vars(args) for args in seen)) + finally: + if original is None: + del sys.modules[name] + else: + sys.modules[name] = original + + def test_complete_literal_cli_compare_and_no_clobber(self): + with tempfile.TemporaryDirectory() as directory: + cwd = Path(directory) + files = [] + for mode in ("static", "dynamic"): + path = cwd / f"{mode}.json" + path.write_text(json.dumps(run_record(mode)), encoding="utf-8") + files.append(path) + output = cwd / "comparison.json" + command = [ + sys.executable, + "-I", + "-B", + str(PUBLIC / "dspark_reproduction.py"), + "compare", + "--runs", + *map(str, files), + "--output", + str(output), + ] + result = subprocess.run(command, cwd=cwd, capture_output=True, text=True, timeout=60) + self.assertEqual(result.returncode, 0, result.stderr) + report = json.loads(output.read_bytes()) + self.assertAlmostEqual(report["dynamic_over_static_pct"], 100 / 3) + self.assertIsNone(report["confidence_interval"]) + self.assertFalse(report["task_accuracy_qualified"]) + self.assertNotEqual( + subprocess.run(command, cwd=cwd, capture_output=True, timeout=60).returncode, 0 + ) + + def test_static_fallback_no_compact_unmatched_and_partial_refused(self): + static, dynamic = run_record("static"), run_record("dynamic") + for field, value, message in ( + ("resolved_confidence_enabled", False, "fell back"), + ("runtime_confidence_enabled", False, "fell back"), + ("routing", {"native_replays": 3, "compact_replays": 0}, "never executed"), + ("run_complete", False, "closed"), + ): + bad = {**dynamic, field: value} + with self.assertRaisesRegex(ValueError, message): + observations.compare_runs([static, bad]) + different = copy.deepcopy(dynamic) + different["comparison_contract"]["physical_k"] = 4 + with self.assertRaisesRegex(ValueError, "differ"): + observations.compare_runs([static, different]) + different = copy.deepcopy(dynamic) + different["responses"][0]["prompt_token_ids"] = [999] + with self.assertRaisesRegex(ValueError, "prompt"): + observations.compare_runs([static, different]) + + def test_internal_constructor_kwargs_not_public_llm_args(self): + # Real argument-validation and constructor controls live in the separate + # native-runtime test. This CPU guard prevents accidental private API injection. + driver = (ROOT / "tensorrt_llm/_torch/speculative/dspark_reproduction.py").read_text() + self.assertNotIn("executor_cls", driver) + self.assertNotIn("runtime_constructors", driver) + + +if __name__ == "__main__": + unittest.main() From 790bd2c19671f5c5d367d6024fcf06ca2a324a01 Mon Sep 17 00:00:00 2001 From: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:38:26 -0700 Subject: [PATCH 5/7] [TRTLLM-14620][fix] Complete public DSpark reproduction integration Preserve artifact reuse and provenance, publish the exercised bounded integration scope, retain generated schemas, and correct the current device-scalar warmup call. Signed-off-by: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> --- .../deepseek_v4/CONFIDENCE_REPRODUCTION.md | 34 ++-- .../models/core/deepseek_v4/COST_EXPORT.md | 29 ++-- .../models/core/deepseek_v4/COST_PROFILING.md | 41 +++-- examples/models/core/deepseek_v4/README.md | 21 ++- setup.py | 6 +- .../engine/runners/decoder/runner.py | 7 +- .../schemas/trtllm-serve-config.schema.json | 162 +++++++++++++++++- .../trtllm-serve-disagg-config.schema.json | 162 +++++++++++++++++- .../test_dspark_model_engine_runtime.py | 57 +++++- 9 files changed, 465 insertions(+), 54 deletions(-) diff --git a/examples/models/core/deepseek_v4/CONFIDENCE_REPRODUCTION.md b/examples/models/core/deepseek_v4/CONFIDENCE_REPRODUCTION.md index 37e277451747..e63a97d7adfc 100644 --- a/examples/models/core/deepseek_v4/CONFIDENCE_REPRODUCTION.md +++ b/examples/models/core/deepseek_v4/CONFIDENCE_REPRODUCTION.md @@ -1,10 +1,21 @@ # Confidence calibration and a fair static/enabled comparison -This is a **source/CPU-reviewed candidate workflow**. These commands have not yet -completed native full-model GPU validation. Do not describe this document as an -already-supported public reproduction recipe until that validation and -publication finish. Its tools have no private harness or cluster-name/path, -Slurm account, fixed K3 or fixed BS512 binding. +These commands have completed a bounded full-model GPU integration for physical +K3 and global BS512 (G64 across eight attention-DP ranks) in an earlier, separately +recorded Python/native environment: whole-step cost profiling/export, genuine +paired calibration collection/fitting, and independent static and enabled-compact +generation with output and natural-shutdown checks. A newer source composition +has separately passed native/API compatibility and a same-source, same-node +static/enabled-compact comparison at physical K3, global BS512 and G64 across +eight attention-DP ranks, with 512 requests of 32 output tokens. Both independent +model lifetimes completed source/native/artifact/input/output/routing, public +comparison and natural-shutdown checks. The original calibration and costs were +reused; exact generated token sequences were not equal. These are distinct +evidence scopes, not fresh current-source cost/calibration measurements, an +exact-source native rebuild, task accuracy, sustained-serving performance or +global clean-install qualification. Original source/native provenance is +preserved. The tools take explicit model, input, K, batch-size and output +arguments, with no private harness, cluster, account or node requirement. The two routes are distinct: @@ -35,10 +46,10 @@ parallelism and original startup/model/shutdown budget. Do not change numerical configuration or disable failed hardware/operator checks to make a command run. The code neither repairs machines nor changes cluster/security policy. -The complete constructor controls run without constructing a model, inside the -actual driver-backed installed runtime before a costly model caller. Some -native imports require an available driver even though these controls do not -execute model or GPU work; do not substitute a driverless import failure: +Optional developer constructor regression checks do not construct a model. +Run them in a matching installed native runtime; some imports still require an +available driver. They are not a required colleague pre-run gate or a +driverless/native GPU qualification: ```bash python -m pytest -q tests/unittest/executor/test_dspark_reproduction_factory.py @@ -69,7 +80,10 @@ Set `--cost-table` and optionally `--calibration` for the dynamic command in step 5. No new cost profiling, forged live fingerprint or admission receipt is needed first. The normal loader still validates table structure, physical K, legal captured graph cells, finite costs and calibration geometry. Missing or -incompatible cost coverage cannot be reported as dynamic throughput. +incompatible cost coverage cannot be reported as dynamic throughput. Default +reuse permits a configured subset of the originally measured graph ladder while +validating the complete original table and fingerprint. Any drain-guard tail +must remain both measured and captured. `validate_confidence_artifact_freshness` stays at its public default `False`. Users who choose `True` must supply the corresponding genuine destination diff --git a/examples/models/core/deepseek_v4/COST_EXPORT.md b/examples/models/core/deepseek_v4/COST_EXPORT.md index 0368d0e645c6..50cd42851041 100644 --- a/examples/models/core/deepseek_v4/COST_EXPORT.md +++ b/examples/models/core/deepseek_v4/COST_EXPORT.md @@ -23,8 +23,11 @@ schema-v2 artifact and can copy its original bytes after structural validation. Both paths use the actual runtime schema validator. Plain graph summaries remain ineligible; that does not prevent reuse of structurally valid supplied costs. -These are source/CPU-prepared development interfaces. The complete new GPU -collection/export workflow has not been exercised or qualified yet. +Direct whole-step collection/export has been exercised in the recorded earlier +physical-K3/G64, global-BS512 environment. Original measurements and fingerprints +are preserved, not relabeled as fresh measurements for later destinations or +sources. See [the reproduction guide](CONFIDENCE_REPRODUCTION.md) for the distinct +validation scopes and remaining limits. For a previously measured cost table and a runtime capturing G16 and G64: @@ -38,9 +41,13 @@ python tensorrt_llm/tools/dspark_cost_export.py \ request count of the actual captured graph, not always the global batch size. With eight attention-DP ranks, BS128 and BS512 correspond to G16 and G64; with attention DP off they do not. Pass the exact graph ladder the engine resolves. -Native V0, positive finite timing values, physical K, cell provenance and graph -coverage must satisfy the existing consumer. Missing cells are not interpolated, -copied from another K or filled with zero. No cluster name or site path is used. +Default reuse permits that configured ladder to be a subset of the originally +measured ladder. The complete original table/fingerprint and native V0, positive +finite timing values, physical K, cell provenance and configured graph coverage +must satisfy the existing consumer. Any drain-guard tail must remain measured +and captured; optional strict freshness still requires the original full ladder. +Missing cells are not interpolated, copied from another K or filled with zero. +No cluster name or site path is used. Add `--output /path/to/new-costs.json` only to deliver a new byte-identical copy. The destination must not exist. No original values, recorded source head or @@ -82,8 +89,9 @@ python -m tensorrt_llm.tools.verify_costs profile \ Use this command under the site's normal MPI/container launcher with the actual GPU count and full startup/generation/export/shutdown budget. No cluster or -account is embedded in the tooling. This source-prepared interface is **not** a -GPU-tested reproduction instruction. No extra sweep is needed before reusing +account is embedded in the tooling. This is the configurable interface used by +the recorded earlier integration, not a promise of portable performance or +universal destination qualification. No extra sweep is needed before reusing valid supplied cost/calibration predictions. Reduce the closed raw observations with the existing `verify_costs.py export` @@ -123,6 +131,7 @@ observations. Changing a label does not extend an existing measured scope. GPU-forward values and overlap `iterLatencyMS` are not replacements for clean complete-loop measurements: the former omits host work, while the latter can span more than one loop. Optional policy-overhead/economic validation stays -separate. Authentic clean-checkout GPU collection, enabled-routing/output checks -and matched benchmarking are still required before publishing this as a verified -end-to-end self-generation recipe. +separate. The recorded earlier whole-step integration does not by itself qualify +a later runtime's enabled routing, outputs or matched benchmark. See +[the reproduction guide](CONFIDENCE_REPRODUCTION.md) for the distinct completed +and pending validation scopes. diff --git a/examples/models/core/deepseek_v4/COST_PROFILING.md b/examples/models/core/deepseek_v4/COST_PROFILING.md index 97a224862e2a..6cd7dccf8462 100644 --- a/examples/models/core/deepseek_v4/COST_PROFILING.md +++ b/examples/models/core/deepseek_v4/COST_PROFILING.md @@ -3,14 +3,15 @@ # Offline DSpark graph-cost profiling -This is a development interface for collecting and exporting measured graph -replay costs. It is not yet a GPU-qualified colleague reproduction recipe. -The default graph-replay collector does not itself generate a serving SPS table, -paired calibration observations, an admission receipt, or an enabled throughput -result. The explicit complete-loop mode and CPU schema-v2 adapter are described -in [COST_EXPORT.md](COST_EXPORT.md); that new producer path remains GPU-unqualified. -The separate CPU calibration fitter below can fit already collected, correctly -paired shards. The complete public clean-checkout workflow is not yet qualified. +This interface collects and exports measured graph replay costs. The direct +whole-step producer/export and separate paired collector/fitter have completed +one recorded physical-K3/G64, global-BS512 GPU integration in an earlier +Python/native environment. See [the reproduction guide](CONFIDENCE_REPRODUCTION.md) +for the distinct newer-source validation scope and remaining limits. +The default graph-replay collector retains its narrower scope: it does not itself +generate a serving SPS table, paired calibration observations, an admission +receipt or an enabled throughput result. The explicit complete-loop mode and +CPU schema-v2 adapter are described in [COST_EXPORT.md](COST_EXPORT.md). The tools contain no private model, dataset, scratch path, historical receipt or source-head pin. Use your own authorized local checkpoint and tokenized inputs. @@ -19,8 +20,9 @@ access-controlled output directory. The Linux collector creates only a new directory with mode0700 and new raw files with mode0600, using exclusive/non-overwriting creation. It does not change existing paths, ACLs, process umask or host settings. Existing/symlink targets are -rejected. POSIX permission behavior still needs actual Linux validation; Windows -CPU controls do not qualify Linux permissions or Windows ACLs. +rejected. The exercised Linux collector used new private output paths; CPU +permission controls or one Linux execution do not qualify Windows ACLs or every +destination filesystem. ## Scope and terminology @@ -118,8 +120,10 @@ The interface can prepare all previously used global batch sizes G16/G32/G64/G128/G192/G256. With attention DP off, G is the full global batch, not BS/8. Smaller points are explicit G inputs too; no pilot batch is hardcoded. Each additional G requires its own captured graph and measured compatible cells. -Preparation does not imply validation of that batch size. Native/GPU validation -is still unperformed for this new tooling, including BS512/1024. +Preparation does not imply validation of every batch size. The recorded K3/G64, +global-BS512 integration has completed; other listed points remain prepared +interfaces unless separately identified as tested. Compatible artifact reuse +does not require a new batch-size sweep. The JSONL input file must contain at least the largest global cohort size, with one object per line and a nonempty list of integer `prompt_token_ids`. Use the @@ -217,12 +221,13 @@ duplicate shards, wrong pairing/width, nonfinite data and invalid prefix order are rejected. Existing output files are never overwritten. This command has CPU math/CLI coverage, including exact agreement with the -established fitter. The new public paired-observation collector is still a -separate integration task: the graph-cost collector does not emit these shards. -No private collector command or environment variable is a public prerequisite -advertised here. CPU fitting does not certify GPU freshness, task accuracy or -positive scheduling economics. Ordinary runtime reuse of a valid supplied STS -vector remains available without fitting it again. +established fitter. The public paired collector and fitter have also completed +the recorded earlier GPU integration. They remain distinct from graph-cost +collection, which does not emit paired shards. No private collector command or +environment variable is a public prerequisite. Neither fitting nor cost +measurement certifies GPU freshness, task accuracy or positive scheduling +economics. Ordinary runtime reuse of a valid supplied STS vector remains +available without fitting it again. ## Next stages and validation boundary diff --git a/examples/models/core/deepseek_v4/README.md b/examples/models/core/deepseek_v4/README.md index fc956dad5651..2a017c684f58 100644 --- a/examples/models/core/deepseek_v4/README.md +++ b/examples/models/core/deepseek_v4/README.md @@ -476,10 +476,15 @@ configuration. Confidence-guided verification is **default-off and experimental**. It proposes the full physical draft block `K`, then may select a smaller, measured verifier-token budget. It does not reduce the -physical drafter length. This draft integration is not a production deployment recipe: strict -receipt generation and performance qualification are external, and the composed source still needs native -build, GPU/full-model, accuracy and performance validation. Optional mixed scheduling has the -separate evidence and route requirements below; no composed-tip GPU qualification is implied. +physical drafter length. This prototype is not a production deployment recipe. Native/API +compatibility and bounded full-model execution have been exercised in the distinct scopes listed +in [the reproduction guide](CONFIDENCE_REPRODUCTION.md). A newer same-source, same-node +static/enabled-compact pair also completed physical-K3/global-BS512/G64, 512-request/32-output-token work with +independent lifetimes and source/output/routing/natural-closure checks. It reused the original +cost/calibration artifacts and did not produce equal token sequences. An exact-source native +rebuild, global clean-install qualification, task accuracy, sustained-serving performance and +optional mixed-route qualification are not implied. Strict receipt generation and performance +qualification remain separate. The source includes candidate [offline cost profiling](COST_PROFILING.md), [cost export and artifact reuse](COST_EXPORT.md), and @@ -491,7 +496,11 @@ remain as stated in those guides; source availability alone is not GPU qualifica `K` is the physical draft length, `G` is the rank-local request count for a captured graph, and `V` is its verifier-token budget. In the cost table, `V=0` denotes the native full-width path; positive `V` values denote measured compact layouts. The table must cover the configured physical -`K` and exact graph ladder. The scheduler never invents a missing cell. +`K` and contain valid native/compact cells for every configured captured graph size. Default reuse +may select a configured subset of the originally measured ladder without rewriting the artifact; +the complete original table and fingerprint are still validated. Optional strict freshness retains +exact-ladder matching, and any drain-guard tail must remain measured and captured. The scheduler +never invents a missing cell. ### Requested configuration versus actual execution @@ -506,7 +515,7 @@ and positive workload-admission checks before constructing engines and workers: | Startup result | Resolved execution | | --- | --- | -| Explicit opt-in with default artifact reuse | Confidence stays enabled; measured costs must still pass schema, finiteness, physical-`K`, exact graph and layout checks. No admission receipt is created or asserted. | +| Explicit opt-in with default artifact reuse | Confidence stays enabled; measured costs must still pass schema, finiteness, physical-`K`, configured graph coverage and layout checks. No admission receipt is created or asserted. | | Strict validation with missing, unreadable, stale or non-positive admission evidence | Both confidence flags become false and all confidence artifact paths/digests are cleared; ordinary feature-off static `K` runs. A warning explains the fallback. | | Physical `K=1` | Feature-off static `K1`; there is no smaller verifier tier. Confidence inputs are ignored and cleared. | | Strict validation with a matching receipt and positive aggregate value | Confidence remains requested; the existing environment, graph, checkpoint and runtime checks must still pass. | diff --git a/setup.py b/setup.py index e22f2bbf07a5..b3d5a7500714 100644 --- a/setup.py +++ b/setup.py @@ -182,11 +182,11 @@ def has_ext_modules(self): 'libs/libtensorrt_llm_ucx_wrapper.so', 'libs/libdecoder_attention_0.so', 'libs/libtensorrt_llm_nixl_wrapper.so', - 'libs/nixl/**/*', + 'libs/nixl/*', 'tensorrt_llm_transfer_agent_binding*.so', 'tensorrt_llm_transfer_agent_binding.pyi', 'libs/libtensorrt_llm_mooncake_wrapper.so', - 'libs/ucx/**/*', + 'libs/ucx/*', 'libs/libpg_utils.so', 'libs/libdecoder_attention_1.so', 'libs/nvshmem/License.txt', @@ -196,7 +196,7 @@ def has_ext_modules(self): 'deep_ep/LICENSE', 'deep_ep/*.py', 'deep_ep_cpp_tllm.*.so', - "include/**/*", + "include/*", 'deep_gemm/LICENSE', 'deep_gemm/include/**/*', 'deep_gemm/*.py', diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 6931a0c36df4..ded8e8c51cd6 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -3401,6 +3401,7 @@ def _warmup_dspark_ragged_compressor_metadata(self) -> None: "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", @@ -3431,6 +3432,7 @@ def _warmup_dspark_ragged_compressor_metadata(self) -> None: 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 @@ -3479,16 +3481,17 @@ def _warmup_dspark_ragged_compressor_metadata(self) -> None: 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, - num_generations, - total_compressed_tokens, + 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_() diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 1211e43a6390..c5b5bc4f3294 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -930,7 +930,7 @@ }, "DSparkDecodingConfig": { "additionalProperties": false, - "description": "Configuration for DSpark speculative decoding.\n\nDSpark (DeepSeek) is a target-dependent, \"semi-parallel\" speculative\ndecoding method. Like DFlash it captures hidden states from several target\nlayers as cross-attention context and drafts a whole block in a single\nbackbone forward, but it additionally refines the per-position draft logits\nwith a lightweight sequential head (a low-rank Markov head, optionally an RNN\nhead). A confidence head can score each draft position for verification\nscheduling without truncating the proposed block. This configuration selects\nnative full-block drafting and does not expose confidence-scheduling or\ncalibration controls.\n\nKey features:\n- Target-dependent: captures hidden states from ``target_layer_ids``.\n- Semi-parallel: one block backbone forward + cheap sequential head refine.\n- Full-block proposals: confidence scoring does not change the draft tokens\n or the standard target-verification acceptance rule.\n\nReference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec).", + "description": "Configuration for DSpark speculative decoding.\n\nDSpark (DeepSeek) is a target-dependent, \"semi-parallel\" speculative\ndecoding method. Like DFlash it captures hidden states from several target\nlayers as cross-attention context and drafts a whole block in a single\nbackbone forward, but it additionally refines the per-position draft logits\nwith a lightweight sequential head (a low-rank Markov head, optionally an RNN\nhead). A confidence head can score each draft position for verification\nscheduling without truncating the proposed block. Confidence scheduling is\nopt-in; the default configuration retains native full-block verification.\n\nKey features:\n- Target-dependent: captures hidden states from ``target_layer_ids``.\n- Semi-parallel: one block backbone forward + cheap sequential head refine.\n- Full-block proposals: confidence scoring does not change the draft tokens\n or the standard target-verification acceptance rule.\n\nReference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec).", "properties": { "acceptance_rate_threshold": { "anyOf": [ @@ -999,6 +999,89 @@ "description": "Number of draft positions produced per block. If None, read from the draft model config (dspark_block_size). Should equal max_draft_len.", "title": "Block Size" }, + "confidence_admission_receipt_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to an authenticated workload-level confidence admission receipt. Required only when validate_confidence_artifact_freshness=True; it must prove positive aggregate compact value after the fixed confidence-path overhead and safety margin are charged.", + "status": "prototype", + "title": "Confidence Admission Receipt Path" + }, + "confidence_admission_receipt_sha256": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Externally authenticated lowercase SHA256 of confidence_admission_receipt_path. This pin must come from the sealed deployment package rather than from the receipt itself.", + "status": "prototype", + "title": "Confidence Admission Receipt Sha256" + }, + "confidence_mixed_evidence": { + "anyOf": [ + { + "$ref": "#/$defs/DSparkMixedEvidenceConfig" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Optional separately admitted physical-K3, eight-rank BREAKABLE mixed scheduling evidence. Does not reuse decode T(G,V) costs. None, unsupported routes, stale ownership, or uncovered mixed geometry retain native mixed verification while admitted decode remains enabled.", + "status": "prototype" + }, + "confidence_sps_live_fingerprint_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to an independently generated JSON fingerprint of the active runtime. Required when validate_confidence_artifact_freshness=True. Otherwise it is optional observed provenance, kept separate from the cost artifact's original metadata.", + "status": "prototype", + "title": "Confidence Sps Live Fingerprint Path" + }, + "confidence_sps_table_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to an authenticated schema-v2 JSON table of measured decode step costs keyed by exact rank-local CUDA-graph shape (G, V).", + "status": "prototype", + "title": "Confidence Sps Table Path" + }, + "confidence_sts_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to a JSON file with per-position sequential-temperature-scaling (STS) values used to calibrate the confidence head. The scheduler consumes a cumulative product of per-position probabilities, so a per-position calibration error compounds geometrically along the block. If None, the raw sigmoid is used (identity calibration).", + "status": "prototype", + "title": "Confidence Sts Path" + }, "decoding_type": { "const": "DSpark", "default": "DSpark", @@ -1021,6 +1104,20 @@ "description": "Developer interface: dynamically adjust draft length based on active batch size in runtime.Maps batch size to draft lengths.For example: draft_len_schedule = {4:4, 8:2, 32:1} - Batch sizes 1-4: use draft_len=4 - Batch sizes 5-8: use draft_len=2 - Batch sizes 9-32: use draft_len=1 - Batch sizes 33+: use draft_len=0 (implicit, speculation disabled). Mutually exclusive with max_concurrency since draft_len_schedule implicitly support max concurrency control.", "title": "Draft Len Schedule" }, + "enable_confidence_scheduling": { + "default": false, + "description": "Score each drafted position with the draft's confidence head and let the verification scheduler size the batch's verify budget from those scores. The block is always drafted in full; only the number of tokens sent to the target for verification changes, so acceptance (and hence output distribution) is unaffected. Requires a checkpoint with trained confidence-head weights. Gains only appear at high concurrency -- at small batch sizes the step cost is nearly flat in the verified token count, so trimming saves nothing and costs acceptance length.", + "status": "prototype", + "title": "Enable Confidence Scheduling", + "type": "boolean" + }, + "enable_fused_confidence_scheduler": { + "default": false, + "description": "Use the optional policy-neutral Triton implementation of confidence top-k allocation plus exact CUDA-graph bucket fill. This changes only the scheduler launch path: budgets, survival thresholds, selected windows, and fallback behavior remain identical. It is independently default-off and requires confidence scheduling.", + "status": "prototype", + "title": "Enable Fused Confidence Scheduler", + "type": "boolean" + }, "enable_penalty": { "default": false, "description": "If true, enables the occurrence penalties (repetition / presence / frequency) for one-model speculative decoding. Off by default because the penalties need a [num_seq_slots, vocab_size] occurrence-count workspace that is allocated up front (CUDA graphs capture fixed buffer addresses). While off, a request that asks for any of these penalties is rejected at admission rather than silently decoded without them.", @@ -1209,11 +1306,74 @@ "status": "prototype", "title": "Use Rejection Sampling", "type": "boolean" + }, + "validate_confidence_artifact_freshness": { + "default": false, + "description": "Require exact live cost-artifact identity and a positive workload admission receipt before enabling confidence scheduling. By default, explicitly enabled scheduling reuses structurally valid cost and calibration artifacts, with a warning that performance predictions may be stale. This never changes target verification or acceptance.", + "status": "prototype", + "title": "Validate Confidence Artifact Freshness", + "type": "boolean" } }, "title": "DSparkDecodingConfig", "type": "object" }, + "DSparkMixedEvidenceConfig": { + "additionalProperties": false, + "description": "Optional measured BREAKABLE mixed evidence, independent of decode costs.\n\nAll SHA256 pins come from the authenticated deployment package. The live\nfingerprint must be independently observed for the current runtime; copying\nthe cost table's identity does not qualify it. Uncovered or invalid mixed\nevidence keeps full-width native verification without disabling admitted\ndecode confidence scheduling.", + "properties": { + "admission_path": { + "description": "Qualified mixed admission receipt JSON.", + "title": "Admission Path", + "type": "string" + }, + "admission_sha256": { + "description": "Receipt SHA256.", + "pattern": "^[0-9a-f]{64}$", + "title": "Admission Sha256", + "type": "string" + }, + "costs_path": { + "description": "Measured mixed geometry/action cost JSON.", + "title": "Costs Path", + "type": "string" + }, + "costs_sha256": { + "description": "Cost-file SHA256.", + "pattern": "^[0-9a-f]{64}$", + "title": "Costs Sha256", + "type": "string" + }, + "live_fingerprint_path": { + "description": "Independently generated mixed runtime fingerprint JSON.", + "title": "Live Fingerprint Path", + "type": "string" + }, + "live_fingerprint_sha256": { + "description": "Live fingerprint SHA256.", + "pattern": "^[0-9a-f]{64}$", + "title": "Live Fingerprint Sha256", + "type": "string" + }, + "protocol_sha256": { + "description": "Measured protocol SHA256.", + "pattern": "^[0-9a-f]{64}$", + "title": "Protocol Sha256", + "type": "string" + } + }, + "required": [ + "costs_path", + "costs_sha256", + "admission_path", + "admission_sha256", + "protocol_sha256", + "live_fingerprint_path", + "live_fingerprint_sha256" + ], + "title": "DSparkMixedEvidenceConfig", + "type": "object" + }, "DecodeCudaGraphConfig": { "additionalProperties": false, "description": "CUDA graph configuration for decode requests.", diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index 15599ad03a1f..d56c50104248 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -930,7 +930,7 @@ }, "DSparkDecodingConfig": { "additionalProperties": false, - "description": "Configuration for DSpark speculative decoding.\n\nDSpark (DeepSeek) is a target-dependent, \"semi-parallel\" speculative\ndecoding method. Like DFlash it captures hidden states from several target\nlayers as cross-attention context and drafts a whole block in a single\nbackbone forward, but it additionally refines the per-position draft logits\nwith a lightweight sequential head (a low-rank Markov head, optionally an RNN\nhead). A confidence head can score each draft position for verification\nscheduling without truncating the proposed block. This configuration selects\nnative full-block drafting and does not expose confidence-scheduling or\ncalibration controls.\n\nKey features:\n- Target-dependent: captures hidden states from ``target_layer_ids``.\n- Semi-parallel: one block backbone forward + cheap sequential head refine.\n- Full-block proposals: confidence scoring does not change the draft tokens\n or the standard target-verification acceptance rule.\n\nReference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec).", + "description": "Configuration for DSpark speculative decoding.\n\nDSpark (DeepSeek) is a target-dependent, \"semi-parallel\" speculative\ndecoding method. Like DFlash it captures hidden states from several target\nlayers as cross-attention context and drafts a whole block in a single\nbackbone forward, but it additionally refines the per-position draft logits\nwith a lightweight sequential head (a low-rank Markov head, optionally an RNN\nhead). A confidence head can score each draft position for verification\nscheduling without truncating the proposed block. Confidence scheduling is\nopt-in; the default configuration retains native full-block verification.\n\nKey features:\n- Target-dependent: captures hidden states from ``target_layer_ids``.\n- Semi-parallel: one block backbone forward + cheap sequential head refine.\n- Full-block proposals: confidence scoring does not change the draft tokens\n or the standard target-verification acceptance rule.\n\nReference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec).", "properties": { "acceptance_rate_threshold": { "anyOf": [ @@ -999,6 +999,89 @@ "description": "Number of draft positions produced per block. If None, read from the draft model config (dspark_block_size). Should equal max_draft_len.", "title": "Block Size" }, + "confidence_admission_receipt_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to an authenticated workload-level confidence admission receipt. Required only when validate_confidence_artifact_freshness=True; it must prove positive aggregate compact value after the fixed confidence-path overhead and safety margin are charged.", + "status": "prototype", + "title": "Confidence Admission Receipt Path" + }, + "confidence_admission_receipt_sha256": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Externally authenticated lowercase SHA256 of confidence_admission_receipt_path. This pin must come from the sealed deployment package rather than from the receipt itself.", + "status": "prototype", + "title": "Confidence Admission Receipt Sha256" + }, + "confidence_mixed_evidence": { + "anyOf": [ + { + "$ref": "#/$defs/DSparkMixedEvidenceConfig" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Optional separately admitted physical-K3, eight-rank BREAKABLE mixed scheduling evidence. Does not reuse decode T(G,V) costs. None, unsupported routes, stale ownership, or uncovered mixed geometry retain native mixed verification while admitted decode remains enabled.", + "status": "prototype" + }, + "confidence_sps_live_fingerprint_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to an independently generated JSON fingerprint of the active runtime. Required when validate_confidence_artifact_freshness=True. Otherwise it is optional observed provenance, kept separate from the cost artifact's original metadata.", + "status": "prototype", + "title": "Confidence Sps Live Fingerprint Path" + }, + "confidence_sps_table_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to an authenticated schema-v2 JSON table of measured decode step costs keyed by exact rank-local CUDA-graph shape (G, V).", + "status": "prototype", + "title": "Confidence Sps Table Path" + }, + "confidence_sts_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to a JSON file with per-position sequential-temperature-scaling (STS) values used to calibrate the confidence head. The scheduler consumes a cumulative product of per-position probabilities, so a per-position calibration error compounds geometrically along the block. If None, the raw sigmoid is used (identity calibration).", + "status": "prototype", + "title": "Confidence Sts Path" + }, "decoding_type": { "const": "DSpark", "default": "DSpark", @@ -1021,6 +1104,20 @@ "description": "Developer interface: dynamically adjust draft length based on active batch size in runtime.Maps batch size to draft lengths.For example: draft_len_schedule = {4:4, 8:2, 32:1} - Batch sizes 1-4: use draft_len=4 - Batch sizes 5-8: use draft_len=2 - Batch sizes 9-32: use draft_len=1 - Batch sizes 33+: use draft_len=0 (implicit, speculation disabled). Mutually exclusive with max_concurrency since draft_len_schedule implicitly support max concurrency control.", "title": "Draft Len Schedule" }, + "enable_confidence_scheduling": { + "default": false, + "description": "Score each drafted position with the draft's confidence head and let the verification scheduler size the batch's verify budget from those scores. The block is always drafted in full; only the number of tokens sent to the target for verification changes, so acceptance (and hence output distribution) is unaffected. Requires a checkpoint with trained confidence-head weights. Gains only appear at high concurrency -- at small batch sizes the step cost is nearly flat in the verified token count, so trimming saves nothing and costs acceptance length.", + "status": "prototype", + "title": "Enable Confidence Scheduling", + "type": "boolean" + }, + "enable_fused_confidence_scheduler": { + "default": false, + "description": "Use the optional policy-neutral Triton implementation of confidence top-k allocation plus exact CUDA-graph bucket fill. This changes only the scheduler launch path: budgets, survival thresholds, selected windows, and fallback behavior remain identical. It is independently default-off and requires confidence scheduling.", + "status": "prototype", + "title": "Enable Fused Confidence Scheduler", + "type": "boolean" + }, "enable_penalty": { "default": false, "description": "If true, enables the occurrence penalties (repetition / presence / frequency) for one-model speculative decoding. Off by default because the penalties need a [num_seq_slots, vocab_size] occurrence-count workspace that is allocated up front (CUDA graphs capture fixed buffer addresses). While off, a request that asks for any of these penalties is rejected at admission rather than silently decoded without them.", @@ -1209,11 +1306,74 @@ "status": "prototype", "title": "Use Rejection Sampling", "type": "boolean" + }, + "validate_confidence_artifact_freshness": { + "default": false, + "description": "Require exact live cost-artifact identity and a positive workload admission receipt before enabling confidence scheduling. By default, explicitly enabled scheduling reuses structurally valid cost and calibration artifacts, with a warning that performance predictions may be stale. This never changes target verification or acceptance.", + "status": "prototype", + "title": "Validate Confidence Artifact Freshness", + "type": "boolean" } }, "title": "DSparkDecodingConfig", "type": "object" }, + "DSparkMixedEvidenceConfig": { + "additionalProperties": false, + "description": "Optional measured BREAKABLE mixed evidence, independent of decode costs.\n\nAll SHA256 pins come from the authenticated deployment package. The live\nfingerprint must be independently observed for the current runtime; copying\nthe cost table's identity does not qualify it. Uncovered or invalid mixed\nevidence keeps full-width native verification without disabling admitted\ndecode confidence scheduling.", + "properties": { + "admission_path": { + "description": "Qualified mixed admission receipt JSON.", + "title": "Admission Path", + "type": "string" + }, + "admission_sha256": { + "description": "Receipt SHA256.", + "pattern": "^[0-9a-f]{64}$", + "title": "Admission Sha256", + "type": "string" + }, + "costs_path": { + "description": "Measured mixed geometry/action cost JSON.", + "title": "Costs Path", + "type": "string" + }, + "costs_sha256": { + "description": "Cost-file SHA256.", + "pattern": "^[0-9a-f]{64}$", + "title": "Costs Sha256", + "type": "string" + }, + "live_fingerprint_path": { + "description": "Independently generated mixed runtime fingerprint JSON.", + "title": "Live Fingerprint Path", + "type": "string" + }, + "live_fingerprint_sha256": { + "description": "Live fingerprint SHA256.", + "pattern": "^[0-9a-f]{64}$", + "title": "Live Fingerprint Sha256", + "type": "string" + }, + "protocol_sha256": { + "description": "Measured protocol SHA256.", + "pattern": "^[0-9a-f]{64}$", + "title": "Protocol Sha256", + "type": "string" + } + }, + "required": [ + "costs_path", + "costs_sha256", + "admission_path", + "admission_sha256", + "protocol_sha256", + "live_fingerprint_path", + "live_fingerprint_sha256" + ], + "title": "DSparkMixedEvidenceConfig", + "type": "object" + }, "DecodeCudaGraphConfig": { "additionalProperties": false, "description": "CUDA graph configuration for decode requests.", 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 7e1b0ffde6a7..94628fa590f8 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 @@ -8,6 +8,10 @@ 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 @@ -152,6 +156,7 @@ def _warmup_metadata(position_helper, mask_helper): 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, ) @@ -164,8 +169,8 @@ def test_ragged_compressor_warmup_primes_and_cleans_intermediate_shape(monkeypat def record_positions(*args): position_calls.append((torch.is_inference_mode_enabled(), args[3:])) - def record_mask(*args): - mask_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( @@ -180,7 +185,8 @@ def record_mask(*args): DecoderRunner._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: 42, 4: 14}, [1, 4]))] + 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() @@ -212,6 +218,51 @@ def fail_positions(*args): 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( + _dspark_confidence_enabled=True, + _dspark_trims_submitted_tokens=True, + attn_metadata=metadata, + _config=SimpleNamespace(max_batch_size=8, cuda_graph_batch_sizes=[1, 4, 8]), + ) + # Exercise the helpers' existing eager branch, not a fake signature or GPU compile. + 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"): + DecoderRunner._warmup_dspark_ragged_compressor_metadata(engine) + else: + DecoderRunner._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(): From 689240df8f4ca3b53d84ba0a43c718c1fd002c35 Mon Sep 17 00:00:00 2001 From: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> Date: Fri, 16 Oct 2026 10:44:15 -0700 Subject: [PATCH 6/7] [TRTLLM-14620][docs] Simplify confidence reproduction workflow Signed-off-by: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> --- .../deepseek_v4/CONFIDENCE_REPRODUCTION.md | 343 +++++++++--------- .../core/deepseek_v4/CONFIDENCE_VALIDATION.md | 244 +++++++++++++ examples/models/core/deepseek_v4/README.md | 228 +----------- 3 files changed, 433 insertions(+), 382 deletions(-) create mode 100644 examples/models/core/deepseek_v4/CONFIDENCE_VALIDATION.md diff --git a/examples/models/core/deepseek_v4/CONFIDENCE_REPRODUCTION.md b/examples/models/core/deepseek_v4/CONFIDENCE_REPRODUCTION.md index e63a97d7adfc..d43acb14eedf 100644 --- a/examples/models/core/deepseek_v4/CONFIDENCE_REPRODUCTION.md +++ b/examples/models/core/deepseek_v4/CONFIDENCE_REPRODUCTION.md @@ -1,97 +1,162 @@ -# Confidence calibration and a fair static/enabled comparison - -These commands have completed a bounded full-model GPU integration for physical -K3 and global BS512 (G64 across eight attention-DP ranks) in an earlier, separately -recorded Python/native environment: whole-step cost profiling/export, genuine -paired calibration collection/fitting, and independent static and enabled-compact -generation with output and natural-shutdown checks. A newer source composition -has separately passed native/API compatibility and a same-source, same-node -static/enabled-compact comparison at physical K3, global BS512 and G64 across -eight attention-DP ranks, with 512 requests of 32 output tokens. Both independent -model lifetimes completed source/native/artifact/input/output/routing, public -comparison and natural-shutdown checks. The original calibration and costs were -reused; exact generated token sequences were not equal. These are distinct -evidence scopes, not fresh current-source cost/calibration measurements, an -exact-source native rebuild, task accuracy, sustained-serving performance or -global clean-install qualification. Original source/native provenance is -preserved. The tools take explicit model, input, K, batch-size and output -arguments, with no private harness, cluster, account or node requirement. - -The two routes are distinct: - -- Reuse a structurally valid supplied calibration/cost table, preserving their - original measurements and provenance. Once confidence is explicitly enabled, - reuse is the default and warnings explain that predictions may be stale. -- Generate new calibration and measured whole-step costs using the offline - commands below. Optional freshness/economic validation is a separate choice; - it is not required before reusing valid artifacts. - -Confidence scheduling remains off by default. Target verification, sampling, -KV/graph contracts, kernel guards and fatal errors are not relaxed by reuse. -Poor cost/calibration estimates can choose a slower verifier shape; reuse is -not a performance guarantee. - -## 1. Build or reuse an ordinary compatible runtime - -Start from a complete public checkout containing these tools, the configurable -cost-profile adapter, and the default-reuse consumer. Use the repository's -[source-build instructions](../../../../docs/source/installation/build-from-source.md) -or its supported precompiled/editable installation route for compatible native -artifacts. Unrelated Python changes do not force native recompilation. Preserve -the original native provenance; the tools separately record loaded Python bytes. - -Use your cluster's ordinary allocation/container/MPI launch procedure and model -mounts. Keep the intended precision, attention backend, fused/prenorm paths, -parallelism and original startup/model/shutdown budget. Do not change numerical -configuration or disable failed hardware/operator checks to make a command run. -The code neither repairs machines nor changes cluster/security policy. - -Optional developer constructor regression checks do not construct a model. -Run them in a matching installed native runtime; some imports still require an -available driver. They are not a required colleague pre-run gate or a -driverless/native GPU qualification: + + + +# DSpark confidence scheduling and reproduction + +Use this guide to run confidence scheduling with existing artifacts or generate +your own. Ordinary inference and benchmarking use the normal TensorRT LLM APIs, +runtime YAML and CLI; no reproduction wrapper or particular cluster is required. + +Confidence scheduling is experimental and **off by default**. When enabled, it +still proposes the full draft block and uses normal target verification and +sampling; it may choose a smaller measured verification layout. + +## 1. Choose your inputs and runtime + +Start from the [DeepSeek-V4 Quick Start](README.md#quick-start) and the ordinary +[installation/build instructions](../../../../docs/source/installation/build-from-source.md). +Use a compatible installed runtime containing these tools; an unrelated Python +change does not by itself require rebuilding native TensorRT LLM. + +You need a supported embedded DSpark checkpoint with **trained confidence-head +weights**, enough GPU/KV capacity, and a workload prepared with its tokenizer. +Temperature calibration does not train or replace those weights. + +Choose these inputs once and use them consistently: + +| Input | Meaning | +| --- | --- | +| Physical `K` | `speculative_config.max_draft_len` in YAML; `--physical-k` in offline/helper commands. It must match the checkpoint/runtime. Dynamic compact scheduling needs K > 1; K1 resolves to static. | +| Global `BS` (offline/helper) | `--batch-size` selects the helper's fixed request cohort. For that cohort with attention DP, rank-local `G = BS / attention-DP size`; otherwise G is the global batch. | +| Ordinary benchmark load | `--max_batch_size` is rank-local with attention DP; `--concurrency` is global concurrent load and `--num_requests` is total submitted work. These are separate controls, not aliases for helper BS. | +| Graph ladder | `cuda_graph_config.batch_sizes` contains rank-local G values covered by the physical-K-specific cost table. Runtime request/token/KV limits must accommodate them. | +| Model and data | Your checkpoint and workload paths; keep precision, backend, sampling and output work fixed for comparisons. | +| Artifacts and outputs | A supplied or self-generated cost table; optional temperature JSON; fresh writable output paths. Inputs can remain read-only. | + +The interfaces accept the previously used BS128/256/512/1024/1536/2048 grid and +smaller points when supported by the actual graph/K/cost/capacity configuration. +This is not a validation claim for every point. The example below uses K3 and +G64 on eight attention-DP ranks; adjust the declared inputs for your workload. +Other route restrictions are in the [reference](CONFIDENCE_VALIDATION.md#supported-routes-and-configuration). + +## 2. Reuse artifacts and enable confidence in normal YAML + +Start with your working DeepSeek-V4 runtime YAML. Keep its model, topology, +precision, KV and backend settings. For this example, use padded graphs and +replace its speculative configuration with: + +```yaml +cuda_graph_config: + enable_padding: true + batch_sizes: [64] +speculative_config: + decoding_type: DSpark + max_draft_len: 3 + enable_confidence_scheduling: true + confidence_sps_table_path: /path/to/artifacts/costs.json +``` + +Save this as `dynamic.yaml`. For `static.yaml`, keep the same numerical/runtime +settings but set `enable_confidence_scheduling: false` and **omit all +confidence-only artifact, freshness, mixed-evidence and fused-scheduler fields**. +Do not replace the checkpoint's full model configuration with this YAML fragment. + +The cost table is required for dynamic scheduling. It must contain finite +positive native/compact measurements compatible with physical K, captured graph +sizes and layouts. Default reuse allows a configured subset of a measured graph +ladder; any priced drain tail must still be measured and captured. Missing +coverage is an error, not a silently timed static result. + +Temperature calibration is optional. Omit `confidence_sts_path` to use raw +confidence probabilities, equivalent to unit temperatures. To reuse a fitted +vector, add `confidence_sts_path: /path/to/artifacts/calibration.json` under +`speculative_config`. + +By default, structurally valid artifacts can be reused across source revisions +and nodes with warnings that predictions may be stale. Their original provenance +is retained; different kernels or hardware can change prediction quality and +performance. New profiling is your choice, not a prerequisite to every run. +Strict freshness/workload-admission validation is an +[optional advanced mode](CONFIDENCE_VALIDATION.md#optional-strict-freshness-and-admission). + +## 3. Run ordinary inference or benchmarks + +Pass `dynamic.yaml` through the normal `LLM` configuration or +`trtllm-serve --config` interface described in [Serving](README.md#serving). +No calibration-collection/comparison helper is required for ordinary use. + +For benchmarking, follow [DeepSeek-V4 Benchmarking](README.md#benchmarking) and +the [benchmarking guide](../../../../docs/source/developer-guide/perf-benchmarking.md#preparing-a-dataset) +to prepare a standard benchmark dataset. It uses `task_id`, `input_ids` or +`prompt`, and `output_tokens` on each JSONL line. + +For an explicit greedy, fixed-output comparison, save this as `sampler.yaml`: + +```yaml +temperature: 0.0 +top_k: 1 +ignore_eos: true +``` + +Run each arm in its own process/model lifetime, changing only the configuration +and report paths: ```bash -python -m pytest -q tests/unittest/executor/test_dspark_reproduction_factory.py +trtllm-bench --model --model_path /path/to/checkpoint throughput \ + --tp 8 --ep 8 --dataset /path/to/benchmark.jsonl \ + --max_batch_size 64 --max_num_tokens 8192 --max_seq_len 4096 \ + --concurrency 512 --num_requests 512 --warmup 512 \ + --sampler_options /path/to/sampler.yaml --config /path/to/static.yaml \ + --report_json /path/to/results/static.json + +trtllm-bench --model --model_path /path/to/checkpoint throughput \ + --tp 8 --ep 8 --dataset /path/to/benchmark.jsonl \ + --max_batch_size 64 --max_num_tokens 8192 --max_seq_len 4096 \ + --concurrency 512 --num_requests 512 --warmup 512 \ + --sampler_options /path/to/sampler.yaml --config /path/to/dynamic.yaml \ + --report_json /path/to/results/dynamic.json ``` -This exercises the actual public argument validator and the existing internal -factory seam without constructing a model. Internal `executor_cls` is not a -public LLM argument and remains rejected. Passing this check is not a native -GPU/operator/model result. - -Prepare tokenized JSONL inputs: each line has a nonempty integer -`prompt_token_ids` list. Model weights, inputs and raw observations stay in -access-controlled storage; they are not published with the tooling. The selected -first `--batch-size` rows form one fixed request cohort. Each request generates -the same requested output length with greedy sampling and EOS ignored. The -original sequence-length limit must contain prompt plus output. - -Use a YAML accepted by `TorchLlmArgs`: matching physical K, intended topology, -enough rank-local request/token/KV capacity and padded decode graphs. Aggregate -PP=1/CP=1 generation is supported by this initial driver. Guided decoding, -LoRA, custom logits processing, disaggregated transfer, in-graph sampling and -resource-governor control are not part of this bounded fixed-cohort workflow. -No alternative attention backend is selected by the tool. - -## 2. Reuse supplied artifacts, if that is your choice - -Set `--cost-table` and optionally `--calibration` for the dynamic command in -step 5. No new cost profiling, forged live fingerprint or admission receipt is -needed first. The normal loader still validates table structure, physical K, -legal captured graph cells, finite costs and calibration geometry. Missing or -incompatible cost coverage cannot be reported as dynamic throughput. Default -reuse permits a configured subset of the originally measured graph ladder while -validating the complete original table and fingerprint. Any drain-guard tail -must remain both measured and captured. - -`validate_confidence_artifact_freshness` stays at its public default `False`. -Users who choose `True` must supply the corresponding genuine destination -fingerprint and workload-admission artifacts; this driver does not fabricate -them. Original cost/calibration hashes are recorded separately from the runtime -actually used for a new comparison. - -## 3. Generate calibration from real native full-K proposals +Both YAML files must include your intended `enable_attention_dp: true` and +matching topology for this example. With attention DP, `--max_batch_size` is +rank-local whereas `--concurrency` is global. Match the dataset output lengths +and ensure prompt plus output fits the sequence limit. Explicit CLI settings +override YAML values. Use your site's normal allocation/container procedure; +do not launch multiple competing copies on the same GPUs. + +This ordinary confidence benchmark recipe has been checked against the public +source interfaces, **not GPU-validated as a new recipe**. Standard benchmark +reports use their own load/timing protocol and do not establish compact routing +by themselves. For a controlled routing-checked comparison, use the optional +helper below. Neither protocol is interchangeable with historical plateau TPS. + +## 4. Generate new artifacts when useful + +You can regenerate either artifact independently. Confidence scheduling does not +require temperature fitting before cost profiling or benchmarking. + +### Whole-step costs + +Follow [Offline cost profiling](COST_PROFILING.md) and +[Cost export](COST_EXPORT.md). Declare native V0 and legal compact cells at the +chosen K and rank-local graph ladder, enable `enable_iter_perf_stats: true` in +the profiling YAML, and collect with `--cost-scope whole-step`. Reduce the +completed all-rank observations, then export the schema-v2 runtime table with +`dspark_cost_export.py`. + +Use directly measured complete-loop costs; graph-replay medians are diagnostics, +not serving prices. The current whole-step exporter supports attention-DP fanout +or a single rank, not multi-rank attention-DP-off export. The linked guides give +the grid, sample floors, output files and exact commands. + +### Optional temperature calibration + +The paired collector uses a separate tokenized input format: +each JSONL line has a nonempty integer `prompt_token_ids` list. Set the shell +variables below to the inputs chosen in step 1; `RUNTIME_YAML` is the matching +runtime YAML with overlap scheduling enabled. `SOURCE_ROOT` is your checkout, +`RUN_ROOT` is private writable storage, and the timeout reserves the full +startup/model/shutdown budget. ```bash mpirun -n "$RANKS" python tensorrt_llm/tools/dspark_reproduction.py run \ @@ -105,56 +170,19 @@ python tensorrt_llm/tools/dspark_calibration.py \ --physical-k "$K" --output "$RUN_ROOT/calibration.json" ``` -The current collector supports the original overlap scheduler input seam; -`disable_overlap_scheduler: False` must already be your intended configuration. -It runs native full-K verification and evaluates the existing confidence head. -There is no confidence-selected shortening and no loaded cost table. Real -warm-up cohorts may contribute calibration observations, but never benchmark -timing. The fitter's default sample floor is 1000 genuine paired rows; collect -enough requests/output tokens rather than weakening that floor for real use. - -Pairing is not by arrival order. Step n's confidence and newly proposed tokens -describe the proposal verified later, usually at n+1. The adapter binds that -later forward to the exact original input `SampleState.device` object. It -snapshots verification-input proposals before persistent buffers can be reused, -and captures the producing confidence row and draft stamp before the normal -sampler records its existing event. It reads those host copies only after the -original sampler's normal synchronization/update. Request incarnation, source -proposal content, worker-local confidence row, sampler slot and full-K extent -must agree. Prefill placeholders, retired requests and unscored bootstrap rows -are counted, not fitted as genuine observations. Wrong-step/stale/censored -pairs fail; the current live confidence is never substituted. - -Expected files are `rank-N-pairs.pt`, `rank-N-run.json` and the fitted JSON with -`sts_temperatures`/`temperatures`. Each shard has logits `[rows,K]`, cumulative -accepted-prefix labels, exact join metadata and actual source/runtime identity. -These observations do not establish downstream task accuracy by themselves. - -## 4. Generate directly measured costs - -Follow [Offline cost profiling](COST_PROFILING.md) and -[Whole-step export](COST_EXPORT.md). Declare native V0 plus legal compact cells -for your configurable physical K and resolved rank-local graph tiers. Use the -`--cost-scope whole-step` producer and export its actual completed measurements -to the schema-v2 table consumed by `confidence_sps_table_path`. - -Do not substitute graph-only timing, sum overlapping components or invent -serving prices from calibration samples. The whole-step producer reuses the -original executor-loop timing boundary; profile-only forced cells cannot serve -ordinary traffic or prove automatic-policy gains. Keep measured confidence-head -and policy cost boundaries explicit. This candidate benchmark measures actual -online policy in its E2E time; it does not manufacture a separate overhead or -economic-admission receipt from that result. - -For attention DP on, global BS determines rank-local G through the attention-DP -fanout. The current schema-v2 cost exporter supports this fanout or one rank; -multi-rank attention-DP-off costs need their separate schema decision. This -restriction belongs to this cost export, not to generic DSpark support. +The collector runs native full-K verification and joins real confidence logits +to the later accepted-prefix labels for the same proposal and request lifetime. +The fitter requires at least 1000 genuine paired rows by default and produces +per-position temperatures. Collection/fitting is not head training or task +accuracy evaluation. See the [pairing and output contract](CONFIDENCE_VALIDATION.md#paired-collection-and-developer-checks). -## 5. Benchmark ordinary static and genuinely enabled execution +## 5. Optional controlled comparison and routing checks -Each command is an independent model lifetime. Use the same original inputs, -physical K, runtime YAML, output work, warm-up count and full feasible budget: +`dspark_reproduction.py` is for paired collection and controlled fixed-cohort +experiments. It is not a required serving launcher. Its `compare` command +consumes its **own** `rank-N-run.json` records, not `trtllm-bench` reports. +To use it, run its own independent static and dynamic lifetimes with the same +inputs, configuration, output work, warm-up count and budget: ```bash mpirun -n "$RANKS" python tensorrt_llm/tools/dspark_reproduction.py run \ @@ -167,7 +195,7 @@ mpirun -n "$RANKS" python tensorrt_llm/tools/dspark_reproduction.py run \ --mode dynamic --model "$MODEL_DIR" --config "$RUNTIME_YAML" \ --inputs "$INPUTS_JSONL" --source-root "$SOURCE_ROOT" \ --physical-k "$K" --batch-size "$BS" --output-tokens "$OUTPUT_TOKENS" \ - --cost-table "$COST_TABLE_JSON" --calibration "$CALIBRATION_JSON" \ + --cost-table "$COST_TABLE_JSON" \ --timeout "$FULL_BUDGET_SECONDS" --output "$RUN_ROOT/dynamic" python tensorrt_llm/tools/dspark_reproduction.py compare \ @@ -175,27 +203,16 @@ python tensorrt_llm/tools/dspark_reproduction.py compare \ --output "$RUN_ROOT/comparison.json" ``` -This uses the normal worker/request/executor/automatic confidence policy, not -the prescribed profile-only engine. Benchmark arms make no calibration copies. -Actual completed decode graph keys are recorded. Warm-up routing is reset before -the retained cohort; requested flags alone do not prove enabled execution. -The comparison requires both resolved config and runtime enabled status, actual -compact replay in the retained dynamic cohort, unchanged numerical runtime/input -work, completed outputs and natural shutdown. An enabled policy that only chose -native is a valid execution, but this consumer refuses to label it a reproduced -compact/dynamic gain. - -Reported TPS is global fixed-cohort enqueue-to-final-response throughput, -including prefill and final drain. It is not measured TPOT, sustained-refill -throughput, isolated verifier time or the historical plateau measurement. -No confidence interval is fabricated. Repeat in balanced order with fresh -lifetimes if you need a stable estimate; feed equal numbers of static and dynamic -records to `compare`. Exact output-token equality is reported separately and is -not silently substituted for the model's original task-accuracy tests. - -All previously tried batch sizes 128/256/512/1024/1536/2048 and source-backed -smaller points use the same CLI, with compatible configured graph/K/cost coverage. -K is an input, not fixed to K3. K1 native calibration/static execution is legal -but cannot demonstrate compact dynamic scheduling. There is no new broad BS -validation sweep implied: existing and freshly tested points must be labeled -truthfully, with unperformed points listed as prepared, not validated. +Calibration is deliberately omitted above; optionally add +`--calibration "$CALIBRATION_JSON"` to the dynamic command. The helper records +resolved/runtime enablement, native/compact replay, source/artifact identity, +outputs and shutdown. An enabled policy may legally choose native throughout; +`compare` refuses to call that a reproduced compact gain. + +Its TPS covers fixed-cohort enqueue through final response, including prefill +and drain. Use balanced fresh-lifetime repetitions for a stable estimate and +equal static/dynamic record counts for `compare`. Token equality is reported +diagnostically, not substituted for task-accuracy tests. + +See [Validation and developer reference](CONFIDENCE_VALIDATION.md) for the +completed evidence scopes, advanced mixed/freshness contracts and optional tests. diff --git a/examples/models/core/deepseek_v4/CONFIDENCE_VALIDATION.md b/examples/models/core/deepseek_v4/CONFIDENCE_VALIDATION.md new file mode 100644 index 000000000000..ef87b5c963f0 --- /dev/null +++ b/examples/models/core/deepseek_v4/CONFIDENCE_VALIDATION.md @@ -0,0 +1,244 @@ + + + +# Confidence scheduling: validation and developer reference + +Start with the [usage guide](CONFIDENCE_REPRODUCTION.md) for artifact reuse, +self-generation and ordinary inference/benchmarking. This reference separates +validation history and advanced contracts from that user workflow. + +## Recorded validation scope + +These are the distinct integration scopes recorded before this documentation +reorganization. Original source/native/artifact provenance remains unchanged. + +| Scope | Recorded result | Does not establish | +| --- | --- | --- | +| Earlier physical-K3, global-BS512/G64 across eight attention-DP ranks, separately recorded Python/native environment | Full-model whole-step profiling/export, genuinely paired calibration collection/fitting, and independent static/enabled-compact generation completed with output and natural-shutdown checks. | Fresh measurements for another source/runtime/node or every other K/BS. | +| Newer source composition, same-source/same-node K3, global-BS512/G64, 512 requests × 32 output tokens | Native/API compatibility and independent static/enabled-compact lifetimes completed source/native/artifact/input/output/routing, public comparison and natural-shutdown checks, reusing original costs/calibration. Exact generated token sequences differed. | Exact-source native rebuild, global clean install, task accuracy or sustained-serving performance. | +| Ordinary `trtllm-bench --config` confidence recipe in the usage guide | Public CLI/YAML/sampler interfaces are source-supported. | A new GPU execution or performance result for that recipe. | +| Small-K drafter-fusion extension | Source-matched extension passed 98 GPU unit cases on B300/SM103. | Full-model/performance qualification of every composed integration, SM100 qualification or arbitrary-K kernel support. | + +Other prepared batch sizes and configurable physical K are interface coverage, +not additional measured results. Existing source tests and historical +measurements must not be relabeled as current-head GPU validation. + +The controlled helper reports fixed-cohort enqueue-to-final-response TPS, +including prefill and drain. Historical plateau measurements and ordinary +benchmark load protocols use different timing boundaries; percentages from them +must not be pooled. Exact output-token equality is a diagnostic field, not a +replacement for original task-accuracy tests or a universal bitwise requirement. +No confidence interval or best-static/best-dynamic claim is produced by one pair. + +## Supported routes and configuration + +The confidence route currently supports embedded DeepSeek-V4 DSpark with trained +confidence-head weights, PP=1 and padded CUDA graphs. Standalone drafters retain +native verification; disable confidence scheduling for that route. The ordinary +target verification, sampling/acceptance, KV/graph, kernel and fatal-error checks +remain in force. + +Guided decoding, LoRA, per-request logits post-processors and incompatible sparse +top-k paths are unsupported. In particular, `enable_heuristic_topk` and +`use_cute_dsl_topk` assume uniform speculative strides. On ragged steps, +`use_cute_dsl_paged_mqa_logits` falls back to the expanded DeepGEMM path. +Advanced rejection draft sampling is unsupported because its sampled prefix +differs from the greedy prefix that produced the confidence scores. Greedy +requests, including `top_k=1`, and requests with rejection draft sampling +disabled retain their supported behavior. Confidence-off execution retains its +ordinary sampling/post-processing paths. + +The offline collector and fixed-cohort helper additionally require aggregate +PP=1/CP=1 generation, embedded DSpark, MPI world matching TP, sufficient +rank-local request/token/KV capacity and padded decode graphs. They reject +guided decoding, LoRA, custom logits processing, disaggregated KV transfer, +in-graph sampling, encode-only and resource-governor control. Paired collection +uses the original overlap-scheduler seam. These helper restrictions are not a +claim that all ordinary TensorRT LLM routes share the same limitations. + +### Cost/configuration contract + +`V=0` denotes native full-width verification; positive V denotes a measured +compact layout at rank-local graph size G. Physical K is not reduced by +scheduling. K1 has no compact tier and resolves to feature-off static execution. +Equal row counts with different graph/attention layouts are not interchangeable. + +The schema-v2 table must contain finite positive measurements and compatible +physical-K, graph/layout and capacity coverage. Default reuse can select a +configured subset of the original measured ladder while validating the entire +original table/fingerprint; a drain-guard tail must still be measured and +captured. Strict freshness requires the exact original ladder. Missing cells +are not interpolated, copied from another K or fabricated. + +The current whole-step exporter supports attention-DP fanout or one rank; +multi-rank attention-DP-off export is refused. This is an export-schema limit, +not a blanket DSpark runtime restriction. Replay-entry timing remains a separate +diagnostic; it is not renamed to a complete-loop cost or summed with overlapping +host/GPU spans. + +Optional temperature vectors must have finite positive values for each physical +draft position. Omitting calibration uses raw sigmoid probabilities. Reuse +retains original artifacts and records the actual destination runtime separately; +it does not rewrite old fingerprints or guarantee positive performance. + +### Requested configuration versus execution + +| Case | Resolved behavior | +| --- | --- | +| Explicit confidence enablement, default artifact reuse | Enabled, subject to structural/runtime checks; warnings state freshness/economics are unchecked. | +| Optional strict freshness with missing, stale, unreadable or non-positive evidence | Feature-off static; confidence flags and artifact paths/digests are cleared with a warning. | +| Physical K1 | Feature-off static K1, confidence inputs cleared. | +| Strict freshness with matching positive evidence | Remains requested; normal runtime/checkpoint/graph checks still apply. | + +Startup feature-off fallback is not an enabled-native policy decision. An enabled +policy can choose native while retaining its confidence/selection work. Optional +fused-scheduler failure can instead fall back to the tensor scheduler for an +affected graph size. Record resolved and runtime enablement plus actual routing; +a YAML flag or successful server start is insufficient. + +The selector retains the table's immediate predicted-goodput margin, normally +1%, and any supplied iteration/drain guard. Default reuse without a drain guard +uses immediate goodput; strict validation without one retains native +verification. Fixed confidence-path overhead is a separate workload-level +charge, not redistributed to individual candidate choices. None certifies an +end-to-end speedup. + +### Drafter fusion + +The BF16 CuTe attention and RMSNorm–RoPE path is eligible for physical K=1..8 +under the existing SM100/SM103 geometry/layout guards; unsupported inputs retain +the reference path. The eight-row storage bound is unchanged. This dispatch is +independent of the confidence flag and therefore also affects static DSpark. +See the validation table above for the actual exercised scope. + +## Optional strict freshness and admission + +`validate_confidence_artifact_freshness: true` opts into the stricter path. +It consumes four externally supplied inputs: + +- `confidence_sps_table_path`: schema-v2 measured costs. +- `confidence_sps_live_fingerprint_path`: independently observed destination + fingerprint, not a relabeled historical fingerprint. +- `confidence_admission_receipt_path`: schema-v1 workload receipt. +- `confidence_admission_receipt_sha256`: externally trusted SHA256 of its bytes. + +Ordinary artifact reuse does not require these additional artifacts. The public +cost/calibration commands do not produce a positive strict-admission receipt. +A separately reviewed producer must retain measured costs, policy overhead, +calibration, workload, selector-replay and runtime evidence. Source, workload or +artifact changes require reassessing that evidence's applicability. + +The receipt envelope has exactly `schema_version: 1`, `admission` and +`admission_sha256`. The last field hashes the canonical admission body: +`json.dumps(body, sort_keys=True, separators=(",", ":")).encode()`. +Its fields are: + +```text +physical_k, source_head, source_diff_sha256, runtime_snapshot +sps_cost_table_sha256, live_engine_fingerprint_sha256, selector_identity_sha256 +workload_identity_sha256, calibration_result_sha256, selector_replay_sha256 +policy_steps, compact_choices, gross_compact_value_ms_lower_bound +fixed_confidence_path_overhead_ms_upper_bound, safety_margin_ms, admitted +``` + +The selector identity must match `EXACT_SPS_SELECTOR_IDENTITY_SHA256` in +`tensorrt_llm/_torch/speculative/dspark_planner.py`. Admission recomputes: + +```text +net_value_lower_bound_ms = + gross_compact_value_ms_lower_bound + - policy_steps * fixed_confidence_path_overhead_ms_upper_bound + - safety_margin_ms +admitted = compact_choices > 0 and net_value_lower_bound_ms > 0 +``` + +All policy steps, including native choices, pay enabled-path overhead. Do not +allocate that fixed charge back to candidates using a +`policy_steps / compact_choices` ratio. Admission does not replace per-decision +selection or its original drain/price guards. + +Hashes check integrity/cross-artifact consistency, not signatures or origin. +The consumer does not independently derive running checkpoint/workload identity, +rerun calibration/replay or verify every referenced result file. The external +producer remains the trust boundary; hashing arbitrary JSON or fixtures is not +admission evidence. + +## Optional mixed context/generation scheduling + +`confidence_mixed_evidence` defaults to `None`: mixed steps use native +verification while supported decode-only scheduling remains enabled. Decode +T(G,V) costs do not price mixed work; context chunks/positions, history, +bootstrap rows, ordered ranks and shared body capacity enter the mixed key. + +The evidence object has `costs_path`, `costs_sha256`, `admission_path`, +`admission_sha256`, `protocol_sha256`, `live_fingerprint_path` and +`live_fingerprint_sha256`. Cost and receipt schemas are +`dspark-mixed-breakable-costs-v1` and `dspark-mixed-breakable-admission-v1`. +Pins come from a trusted producer; the live source/runtime/model/GPU/topology/ +selector fingerprint must be observed independently. Decode artifact reuse does +not relax this separate contract. + +The supported mixed route is embedded physical-K3 DSpark, TP8/PP1/CP1, +attention DP, greedy sampling and captured BREAKABLE buckets +`[96, 128, 192, 256, 384, 512]`. Unsupported/uncovered routes stay native for +mixed steps without disabling valid decode scheduling. First context chunks +stay native because prefix-cache resolution can change geometry. Later chunks +need their full-width native rows to fit; quoted context work is not trimmed. + +When a peer prefills, ranks use shared BREAKABLE padding, not a rank-local full +decode graph. Only authenticated generation owners can receive shorter +prefixes; context/bootstrap/padding rows retain required rows. Host-authoritative +lengths are published before attention metadata preparation, without changing +backend, sampler or acceptance. Stale/unready owners remain full-width. + +Selection requires an exact measured geometry/action cell and at least 1% +predicted-goodput margin. Current `rank0_forward_pricing_proxy` prices forward +work, not the whole distributed step. Workload admission must charge CPU curve +preparation, producer-stamp transfer and the existing allgather's fixed trailer +on all policy steps, including native/all-decode decisions. Policy equivalence +does not establish zero overhead. + +There is no public mixed receipt-generation command here. Valid-shaped fixture +JSON is not deployment evidence. A separately reviewed producer must supply +measured costs, accuracy/runtime results, selector replay and positive workload +economics; exact integration GPU/full-model and production evidence remain +separate from CPU contracts. + +## Paired collection and developer checks + +Calibration pairs are not matched by arrival order. Step n confidence describes +a proposal verified later, usually at n+1. The adapter binds that forward to its +original `SampleState.device`, snapshots proposals before buffer reuse, and +records confidence/draft stamps before the normal sampler event. Host copies +are read only after the sampler's existing synchronization/update. + +Request incarnation, proposal content, local confidence row, sampler slot and +full-K extent must agree. Prefill placeholders, retired/unscored bootstrap rows +are counted rather than fitted. Wrong-step/stale/censored observations fail; +current-step logits cannot substitute for the producing logits. Compact +verification cannot label an unverified suffix. + +Collection emits `rank-N-pairs.pt` and `rank-N-run.json`. Tensor-only shards +contain logits [rows,K], cumulative accepted-prefix labels, pairing metadata and +source/runtime identity. The fitter retains shard hashes, sample count, +temperatures and before/after ECE; degenerate positions are reported rather than +promoted as informative fits. This is not task-accuracy qualification. + +The optional constructor test below uses the real argument validator and +internal factory seam without constructing a model. Run it only in a compatible +installed native runtime; imports may require a driver. It is a developer +regression, not a colleague pre-run gate or GPU qualification: + +```bash +python -m pytest -q tests/unittest/executor/test_dspark_reproduction_factory.py +``` + +`executor_cls` is internal and remains rejected as a public LLM argument. +Additional synthetic contract regressions are: + +- `tests/unittest/_torch/speculative/hw_agnostic/test_dspark_workload_admission.py` +- `tests/unittest/_torch/speculative/hw_agnostic/test_dspark_admission_resolution.py` +- `tests/unittest/_torch/speculative/hw_agnostic/test_dspark_artifact_reuse.py` +- `tests/unittest/llmapi/test_dspark_workload_admission_config.py` +- `tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_evidence.py` +- `tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_runtime.py` diff --git a/examples/models/core/deepseek_v4/README.md b/examples/models/core/deepseek_v4/README.md index 2a017c684f58..9cec1eca0463 100644 --- a/examples/models/core/deepseek_v4/README.md +++ b/examples/models/core/deepseek_v4/README.md @@ -36,6 +36,7 @@ for how to build TensorRT LLM from source and start a TRT-LLM Docker container. - [Sparse attention](#sparse-attention) - [KV cache](#kv-cache) - [Quantized checkpoints](#quantized-checkpoints) + - [DSpark confidence verification (prototype)](#dspark-confidence-verification-prototype) - [Notes and Troubleshooting](#notes-and-troubleshooting) @@ -474,226 +475,15 @@ configuration. ## DSpark confidence verification (prototype) -Confidence-guided verification is **default-off and experimental**. It proposes the full physical -draft block `K`, then may select a smaller, measured verifier-token budget. It does not reduce the -physical drafter length. This prototype is not a production deployment recipe. Native/API -compatibility and bounded full-model execution have been exercised in the distinct scopes listed -in [the reproduction guide](CONFIDENCE_REPRODUCTION.md). A newer same-source, same-node -static/enabled-compact pair also completed physical-K3/global-BS512/G64, 512-request/32-output-token work with -independent lifetimes and source/output/routing/natural-closure checks. It reused the original -cost/calibration artifacts and did not produce equal token sequences. An exact-source native -rebuild, global clean-install qualification, task accuracy, sustained-serving performance and -optional mixed-route qualification are not implied. Strict receipt generation and performance -qualification remain separate. - -The source includes candidate [offline cost profiling](COST_PROFILING.md), -[cost export and artifact reuse](COST_EXPORT.md), and -[calibration collection plus matched static/dynamic commands](CONFIDENCE_REPRODUCTION.md). -These use ordinary runtime configuration and configurable physical K and batch size, without -private harnesses or cluster-specific inputs. Their end-to-end validation and publication status -remain as stated in those guides; source availability alone is not GPU qualification. - -`K` is the physical draft length, `G` is the rank-local request count for a captured graph, and -`V` is its verifier-token budget. In the cost table, `V=0` denotes the native full-width path; -positive `V` values denote measured compact layouts. The table must cover the configured physical -`K` and contain valid native/compact cells for every configured captured graph size. Default reuse -may select a configured subset of the originally measured ladder without rewriting the artifact; -the complete original table and fingerprint are still validated. Optional strict freshness retains -exact-ladder matching, and any drain-guard tail must remain measured and captured. The scheduler -never invents a missing cell. - -### Requested configuration versus actual execution - -Set `DSparkDecodingConfig.enable_confidence_scheduling=True` to opt in. By default, structurally -valid existing costs and calibration can be reused even when their recorded source or GPU-node -metadata differ from the current run. A warning makes clear that predictions may be stale and -that positive performance economics have not been established. Reuse does not change target -verification or the acceptance rule. - -Set `validate_confidence_artifact_freshness=True` to require the previous strict live-fingerprint -and positive workload-admission checks before constructing engines and workers: - -| Startup result | Resolved execution | -| --- | --- | -| Explicit opt-in with default artifact reuse | Confidence stays enabled; measured costs must still pass schema, finiteness, physical-`K`, configured graph coverage and layout checks. No admission receipt is created or asserted. | -| Strict validation with missing, unreadable, stale or non-positive admission evidence | Both confidence flags become false and all confidence artifact paths/digests are cleared; ordinary feature-off static `K` runs. A warning explains the fallback. | -| Physical `K=1` | Feature-off static `K1`; there is no smaller verifier tier. Confidence inputs are ignored and cleared. | -| Strict validation with a matching receipt and positive aggregate value | Confidence remains requested; the existing environment, graph, checkpoint and runtime checks must still pass. | - -An **enabled-native** policy decision is different from startup -feature-off fallback: it retains the enabled control path while verifying the native layout. -An enabled compact decision uses a positive measured budget. Optional fused-scheduler failure -falls back to the tensor scheduler for the affected graph size; that is a third, separate fallback. - -Before reporting an "enabled" experiment, inspect the resolved -`TorchLlmArgs.speculative_config.enable_confidence_scheduling` and record both the requested and -resolved configuration. Require observed native/compact routing and completed outputs, rather -than inferring execution from an input YAML flag or a successful server start. Do not count -startup feature-off fallback as an enabled-native run. With confidence disabled, omit all -confidence-only fields; explicitly supplying them is a configuration error. - -Confidence scheduling requires padded CUDA graphs and PP=1. Guided decoding, LoRA and incompatible -sparse top-k settings remain unsupported. Per-request logits post-processors are rejected before -activation. Advanced rejection draft sampling is also unsupported: its sampled proposal prefix -differs from the greedy prefix that produced the confidence scores. Greedy requests (including -`top_k=1`) and requests with rejection draft sampling disabled retain their existing behavior. -Confidence-disabled execution retains its ordinary sampling and post-processing paths. - -### Reusing costs and calibration - -For default reuse, supply `confidence_sps_table_path` and, when using a calibrated head, -`confidence_sts_path`. Costs must be finite positive schema-v2 measurements with native and compact -cells for every configured graph size; STS temperatures must be a finite positive vector with -one value per physical draft position. Invalid artifacts raise errors rather than being silently -timed as enabled static execution. Their original metadata is retained; the runner records its -observed runtime separately instead of rewriting the artifact's provenance. - -For example, the confidence portion of a padded-graph configuration is: +DSpark can use a trained confidence head to choose a measured verifier-token budget while +continuing to propose the full physical draft block. Confidence scheduling is experimental and +off by default. Enabling it requires a compatible cost table; optional temperature calibration +can refine the predictions, but is not required to run. -```yaml -speculative_config: - decoding_type: DSpark - max_draft_len: 3 - enable_confidence_scheduling: true - confidence_sps_table_path: /your/artifacts/k3-costs.json - confidence_sts_path: /your/artifacts/k3-sts.json -``` - -This is a configuration example, not an exact-tip GPU-validated reproduction recipe. Keep the -same checkpoint, physical `K`, inputs, parallelism, graph ladder and other runtime options when -comparing it with static DSpark (`enable_confidence_scheduling: false`, without confidence-only -fields). A fair measured comparison keeps source, configuration, inputs and output work matched, -uses independent model lifetimes, checks actual routing and completed outputs, and records actual -timings. Reused predictions do not invalidate that measurement: retain their original provenance -and label freshness unchecked, without promising the earlier gain. Fresh profiling is optional -and can improve prediction quality. Changing batch size requires a graph ladder covered by the -supplied table, not a separate per-batch-size package. - -Both modes retain the table's immediate predicted-goodput margin (normally 1%). If the table -supplies an `iteration_drain_guard`, its prices and final-drain policy are unchanged. Without that -optional metadata, default reuse selects by immediate goodput alone; strict validation retains -native verification. Neither behavior certifies an end-to-end gain. - -### Optional freshness and workload-admission validation - -With `validate_confidence_artifact_freshness: true`, provide all four prototype inputs: - -- `confidence_sps_table_path`: schema-v2 measured costs with exact physical `K`, rank-local graph - sizes `G`, native `V=0` cells and supported positive verifier budgets. -- `confidence_sps_live_fingerprint_path`: an independently supplied fingerprint for that runtime - and graph ladder, not a copy relabeled from an older cost artifact. -- `confidence_admission_receipt_path`: the schema-v1 workload receipt. -- `confidence_admission_receipt_sha256`: the expected lowercase SHA256 of the **receipt file - bytes**, obtained from a trusted, separately reviewed deployment package. - -The candidate [self-generation workflow](CONFIDENCE_REPRODUCTION.md) provides separate cost -profiling/export and correctly paired calibration collection/fitting commands. They do not -generate a positive strict-admission receipt. The fields above consume artifacts; strict -freshness validation is optional and remains separate from default artifact reuse. Do not -manufacture costs or an admitting receipt from unit-test fixtures. A strict receipt's producer -must preserve the measured cost, -workload, calibration, selector-replay and runtime evidence and review their applicability to the -actual deployment. A source, calibration, physical-`K` or workload change requires reassessment; -historical scores or timing results cannot be relabeled as new-source qualification. - -The receipt envelope contains exactly `schema_version: 1`, `admission` and `admission_sha256`. -The last field hashes the canonical admission-body JSON -(`json.dumps(body, sort_keys=True, separators=(",", ":")).encode()`). -The body contains: - -- `physical_k`, `source_head`, `source_diff_sha256`, `runtime_snapshot`; -- `sps_cost_table_sha256`, `live_engine_fingerprint_sha256`, `selector_identity_sha256`; -- `workload_identity_sha256`, `calibration_result_sha256`, `selector_replay_sha256`; -- `policy_steps`, `compact_choices`, `gross_compact_value_ms_lower_bound`, - `fixed_confidence_path_overhead_ms_upper_bound`, `safety_margin_ms`, `admitted`. - -The selector identity must match `EXACT_SPS_SELECTOR_IDENTITY_SHA256` in -`tensorrt_llm/_torch/speculative/dspark_planner.py`. Replay the existing marginal selector without -allocating fixed overhead back to candidates through a `policy_steps / compact_choices` ratio. -Admission recomputes: - -```text -net_value_lower_bound_ms = - gross_compact_value_ms_lower_bound - - policy_steps * fixed_confidence_path_overhead_ms_upper_bound - - safety_margin_ms - -admitted = compact_choices > 0 and net_value_lower_bound_ms > 0 -``` - -Native decisions still pay the enabled confidence-path overhead, so **all** policy steps count. -Strict admission does not replace per-decision selection: exact measured tiers and the existing -iteration/drain guard still determine whether a compact choice is allowed. - -Hashes provide integrity and cross-artifact consistency checks, **not signatures or proof of -origin**. The runtime checks supplied receipt, cost and fingerprint values; it does not independently -derive the running checkpoint/workload identity, re-execute the calibration or selector replay, -or verify their referenced result files. The external producer/deployment package remains the -trust boundary. Hashing an arbitrary receipt and passing that same digest does not establish -trust or qualify an enabled experiment. - -Contract regressions are in `test_dspark_workload_admission.py`, -`test_dspark_admission_resolution.py` (both under -`tests/unittest/_torch/speculative/hw_agnostic/`) and -`tests/unittest/llmapi/test_dspark_workload_admission_config.py`. These fixtures are synthetic -validator tests, not deployment evidence. `test_dspark_artifact_reuse.py` covers optional -freshness, unchanged structural guards and the absent-drain behavior. - -### Optional mixed context/generation scheduling - -`confidence_mixed_evidence` is an optional nested object. Its default is `None`, which leaves -mixed steps at native verification width and does not change enabled decode-only scheduling. -Mixed costs are not decode `T(G,V)` costs: context chunks and their positions, generation history, -bootstrap requests, ordered ranks and the shared body capacity all affect the mixed cost key. - -The object contains `costs_path`, `costs_sha256`, `admission_path`, `admission_sha256`, -`protocol_sha256`, `live_fingerprint_path` and `live_fingerprint_sha256`. Cost, receipt and live -fingerprint digests pin exact file bytes; the protocol digest identifies the measured protocol. -These pins come from a trusted deployment package. The live fingerprint must be observed independently -for the current source, runtime, model, GPU class, topology and selector, not copied from the cost -table. The cost file uses `dspark-mixed-breakable-costs-v1`; its separate receipt uses -`dspark-mixed-breakable-admission-v1`. The reuse setting applies to decode schema-v2 inputs only; -it does not relax this separate mixed-execution contract. - -The supported mixed route is embedded DSpark with physical `K=3`, TP8, PP1, CP1, attention DP, -greedy sampling and captured BREAKABLE body buckets `[96, 128, 192, 256, 384, 512]`. Standalone -draft models, unsupported routes and uncovered mixed geometry retain native mixed verification. -Missing or invalid mixed evidence does not disable an already enabled decode policy. -First context chunks also remain native because KV preparation may still resolve a prefix-cache -hit and change their geometry. Later chunks are eligible only when their full-width native rows -fit the configured token budget, so resource preparation does not trim a quoted context shape. - -When a peer rank prefills, all ranks use the existing shared BREAKABLE padding path rather than -replaying a rank-local full decode graph. Only authenticated generation owners may receive shorter -verification prefixes; context rows, bootstrap requests and graph-padding requests retain their -required rows. The runner publishes those lengths as host-authoritative before attention metadata -preparation, without changing the attention backend, sampler or acceptance rule. Stale or unready -confidence owners retain full physical `K`. - -Selection requires an exact measured geometry/action cell and retains the minimum 1% predicted -goodput gate. The current pricing field is `rank0_forward_pricing_proxy`, not a measurement of the -whole distributed step. Workload admission must charge the opt-in path's overhead on **all** policy -steps, including all-decode and native choices: the optional path adds CPU curve preparation, -producer-stamp transfer and a fixed trailer to the existing policy allgather. Policy equivalence -alone is not evidence of zero overhead or preserved throughput. - -There is no public mixed receipt-generation command in this change. Supplying valid-shaped JSON -or hashing test fixtures does not qualify a deployment. A separately reviewed producer must supply -matching measured costs, runtime/accuracy results, selector replay and positive workload economics. -Exact-tip GPU/full-model validation and production mixed evidence remain required before making a -deployment or performance claim. The CPU contracts are in `test_dspark_mixed_evidence.py` and -`test_dspark_mixed_runtime.py` under `tests/unittest/_torch/speculative/hw_agnostic/`. - -### Drafter fusion scope - -The BF16 CuTe attention and RMSNorm–RoPE path is eligible for physical `K=1..8` under the existing -SM100/SM103, geometry and layout guards; unsupported inputs retain the reference path. This -dispatch is independent of the confidence flag, so it also affects static DSpark. The eight-row -storage bound is unchanged; no arbitrary-`K` kernel extension is claimed. - -The source-matched extension passed 98 GPU unit cases on B300/SM103. That is not full-model or -performance qualification of this composed integration. SM100 and exact-integration GPU/model -validation remain pending; keep this scope experimental until those gates and review complete. +Start with [Confidence scheduling and reproduction](CONFIDENCE_REPRODUCTION.md) for ordinary +runtime configuration, artifact reuse, self-generation and benchmarking. The +[validation and developer reference](CONFIDENCE_VALIDATION.md) records exercised configurations, +advanced restrictions, optional freshness checks and developer tests. ## Notes and Troubleshooting From 7efa4b5101da09fed8c2abe305576ca78309f220 Mon Sep 17 00:00:00 2001 From: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> Date: Fri, 16 Oct 2026 17:20:54 -0700 Subject: [PATCH 7/7] [TRTLLM-14620][docs] Streamline DSpark confidence setup and profiling Signed-off-by: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> --- .../deepseek_v4/CONFIDENCE_REPRODUCTION.md | 245 +++++------ .../core/deepseek_v4/CONFIDENCE_VALIDATION.md | 40 +- .../models/core/deepseek_v4/COST_EXPORT.md | 183 +++----- .../models/core/deepseek_v4/COST_PROFILING.md | 390 ++++++------------ examples/models/core/deepseek_v4/README.md | 8 +- 5 files changed, 350 insertions(+), 516 deletions(-) diff --git a/examples/models/core/deepseek_v4/CONFIDENCE_REPRODUCTION.md b/examples/models/core/deepseek_v4/CONFIDENCE_REPRODUCTION.md index d43acb14eedf..192a597a1d91 100644 --- a/examples/models/core/deepseek_v4/CONFIDENCE_REPRODUCTION.md +++ b/examples/models/core/deepseek_v4/CONFIDENCE_REPRODUCTION.md @@ -3,49 +3,51 @@ # DSpark confidence scheduling and reproduction -Use this guide to run confidence scheduling with existing artifacts or generate -your own. Ordinary inference and benchmarking use the normal TensorRT LLM APIs, -runtime YAML and CLI; no reproduction wrapper or particular cluster is required. +Use this guide to run with supplied artifacts or generate your own. Ordinary +inference and benchmarking use the normal TensorRT LLM runtime YAML, APIs and +CLI; the optional reproduction helper is not a serving prerequisite. -Confidence scheduling is experimental and **off by default**. When enabled, it -still proposes the full draft block and uses normal target verification and -sampling; it may choose a smaller measured verification layout. +Confidence scheduling is experimental and **off by default**. Explicitly enabling +it retains full physical draft proposals and normal target verification and +sampling, while allowing a smaller measured verification layout. -## 1. Choose your inputs and runtime +## 1. Prepare the runtime and inputs -Start from the [DeepSeek-V4 Quick Start](README.md#quick-start) and the ordinary +Follow the [DeepSeek-V4 Quick Start](README.md#quick-start) and the ordinary [installation/build instructions](../../../../docs/source/installation/build-from-source.md). -Use a compatible installed runtime containing these tools; an unrelated Python -change does not by itself require rebuilding native TensorRT LLM. +Use a compatible runtime and a supported embedded DSpark checkpoint containing +trained confidence-head weights. Temperature calibration does not train those +weights. No particular cluster is required. -You need a supported embedded DSpark checkpoint with **trained confidence-head -weights**, enough GPU/KV capacity, and a workload prepared with its tokenizer. -Temperature calibration does not train or replace those weights. - -Choose these inputs once and use them consistently: +Choose the following inputs for your workload: | Input | Meaning | | --- | --- | -| Physical `K` | `speculative_config.max_draft_len` in YAML; `--physical-k` in offline/helper commands. It must match the checkpoint/runtime. Dynamic compact scheduling needs K > 1; K1 resolves to static. | -| Global `BS` (offline/helper) | `--batch-size` selects the helper's fixed request cohort. For that cohort with attention DP, rank-local `G = BS / attention-DP size`; otherwise G is the global batch. | -| Ordinary benchmark load | `--max_batch_size` is rank-local with attention DP; `--concurrency` is global concurrent load and `--num_requests` is total submitted work. These are separate controls, not aliases for helper BS. | -| Graph ladder | `cuda_graph_config.batch_sizes` contains rank-local G values covered by the physical-K-specific cost table. Runtime request/token/KV limits must accommodate them. | -| Model and data | Your checkpoint and workload paths; keep precision, backend, sampling and output work fixed for comparisons. | -| Artifacts and outputs | A supplied or self-generated cost table; optional temperature JSON; fresh writable output paths. Inputs can remain read-only. | +| Physical `K` | YAML `speculative_config.max_draft_len`, and offline/helper `--physical-k`. It must match the loaded checkpoint/runtime. Compact scheduling needs K > 1; K1 resolves to feature-off static execution. | +| Graph size `G` | Rank-local request count in `cuda_graph_config.batch_sizes`. Captured graph layouts and the cost table must cover it; request, token and KV limits must accommodate it. | +| Verifier budget `V` | Total verifier rows in a captured rank-local compact layout: `G <= V < G*(K+1)`. V0 is a sentinel for native full-width `G*(K+1)` verification, not zero tokens. | +| Helper global `BS` | `--batch-size` is a fixed global request cohort, not a graph setting. With attention DP, G = BS / attention-DP size; otherwise G = BS. | +| Ordinary benchmark load | With attention DP, `--max_batch_size` is rank-local. `--concurrency` limits global simultaneous requests; `--num_requests` caps requests read from the dataset, with 0 meaning the dataset length. `--warmup` counts warm-up requests, not iterations or helper cohorts. | +| Model and data | Checkpoint and workload paths, using the checkpoint's tokenizer. Keep topology, precision, backend, sampling and output work fixed when comparing runs. | +| Artifacts | A compatible cost table is required for dynamic scheduling. Temperature calibration is optional. | + +The examples use K3, G64 and eight attention-DP ranks. These are example inputs, +not universal requirements. See [supported routes and configuration](CONFIDENCE_VALIDATION.md#supported-routes-and-configuration) +for runtime and helper limits. -The interfaces accept the previously used BS128/256/512/1024/1536/2048 grid and -smaller points when supported by the actual graph/K/cost/capacity configuration. -This is not a validation claim for every point. The example below uses K3 and -G64 on eight attention-DP ranks; adjust the declared inputs for your workload. -Other route restrictions are in the [reference](CONFIDENCE_VALIDATION.md#supported-routes-and-configuration). +Observations can contain prompt and output tokens. Use a new private writable +output directory; keep inputs read-only if desired. Offline/helper commands +require fresh output paths with existing parent directories and do not overwrite +them. -## 2. Reuse artifacts and enable confidence in normal YAML +## 2. Reuse artifacts in normal runtime YAML -Start with your working DeepSeek-V4 runtime YAML. Keep its model, topology, -precision, KV and backend settings. For this example, use padded graphs and -replace its speculative configuration with: +Start with a complete working DeepSeek-V4 **runtime YAML**. Keep its topology, +precision, KV, backend and capacity settings, and merge this example into it, +replacing its existing speculative configuration: ```yaml +enable_attention_dp: true cuda_graph_config: enable_padding: true batch_sizes: [64] @@ -56,41 +58,40 @@ speculative_config: confidence_sps_table_path: /path/to/artifacts/costs.json ``` -Save this as `dynamic.yaml`. For `static.yaml`, keep the same numerical/runtime -settings but set `enable_confidence_scheduling: false` and **omit all -confidence-only artifact, freshness, mixed-evidence and fused-scheduler fields**. -Do not replace the checkpoint's full model configuration with this YAML fragment. - -The cost table is required for dynamic scheduling. It must contain finite -positive native/compact measurements compatible with physical K, captured graph -sizes and layouts. Default reuse allows a configured subset of a measured graph -ladder; any priced drain tail must still be measured and captured. Missing -coverage is an error, not a silently timed static result. - -Temperature calibration is optional. Omit `confidence_sts_path` to use raw -confidence probabilities, equivalent to unit temperatures. To reuse a fitted -vector, add `confidence_sts_path: /path/to/artifacts/calibration.json` under -`speculative_config`. - -By default, structurally valid artifacts can be reused across source revisions -and nodes with warnings that predictions may be stale. Their original provenance -is retained; different kernels or hardware can change prediction quality and -performance. New profiling is your choice, not a prerequisite to every run. -Strict freshness/workload-admission validation is an +Save the merged runtime configuration as `dynamic.yaml`. This fragment is not a +complete configuration and does not replace the checkpoint's `config.json`. + +Create `static.yaml` from the same runtime configuration, set +`enable_confidence_scheduling: false`, and omit all confidence-only artifact, +freshness, mixed-evidence and fused-scheduler fields. Keep the physical K and +numerical/runtime settings unchanged. + +The cost table must contain finite positive native V0 measurements and the +compact cells it offers, compatible with physical K and graph/layout coverage. +Default reuse permits a configured subset of a measured graph ladder; any priced +drain tail must still be measured and captured. An enabled policy may legally +choose native throughout, but missing/incompatible coverage is not a compact +execution or a reproduced gain. + +Structurally valid artifacts can be reused across source revisions and nodes +with warnings that freshness and performance economics are unchecked. Original +artifact provenance is retained; a different destination can change prediction +quality and performance. New profiling is optional when reusing compatible +artifacts. Strict freshness and economic admission are an [optional advanced mode](CONFIDENCE_VALIDATION.md#optional-strict-freshness-and-admission). ## 3. Run ordinary inference or benchmarks -Pass `dynamic.yaml` through the normal `LLM` configuration or -`trtllm-serve --config` interface described in [Serving](README.md#serving). -No calibration-collection/comparison helper is required for ordinary use. +Use `dynamic.yaml` through the normal `LLM` configuration or +`trtllm-serve --config` interface in [Serving](README.md#serving). For benchmarking, follow [DeepSeek-V4 Benchmarking](README.md#benchmarking) and -the [benchmarking guide](../../../../docs/source/developer-guide/perf-benchmarking.md#preparing-a-dataset) -to prepare a standard benchmark dataset. It uses `task_id`, `input_ids` or -`prompt`, and `output_tokens` on each JSONL line. +the [dataset preparation guide](../../../../docs/source/developer-guide/perf-benchmarking.md#preparing-a-dataset). +The standard benchmark JSONL format uses `task_id`, `input_ids` or `prompt`, and +`output_tokens`. Prepare at least as many requests as you intend to measure; +`--num_requests` does not repeat a shorter dataset. -For an explicit greedy, fixed-output comparison, save this as `sampler.yaml`: +For a greedy, fixed-output comparison, save this as `sampler.yaml`: ```yaml temperature: 0.0 @@ -98,8 +99,8 @@ top_k: 1 ignore_eos: true ``` -Run each arm in its own process/model lifetime, changing only the configuration -and report paths: +Run each arm in a separate process/model lifetime, changing only the config and +report paths: ```bash trtllm-bench --model --model_path /path/to/checkpoint throughput \ @@ -117,50 +118,50 @@ trtllm-bench --model --model_path /path/to/checkpoint throughput \ --report_json /path/to/results/dynamic.json ``` -Both YAML files must include your intended `enable_attention_dp: true` and -matching topology for this example. With attention DP, `--max_batch_size` is -rank-local whereas `--concurrency` is global. Match the dataset output lengths -and ensure prompt plus output fits the sequence limit. Explicit CLI settings -override YAML values. Use your site's normal allocation/container procedure; -do not launch multiple competing copies on the same GPUs. +Both runtime files use attention DP for this example. Match dataset output +lengths and ensure each prompt plus output fits the sequence limit. Explicit CLI +settings override conflicting YAML values. Use the normal allocation/container +procedure; `trtllm-bench` handles process launch, so do not add an MPI wrapper. -This ordinary confidence benchmark recipe has been checked against the public -source interfaces, **not GPU-validated as a new recipe**. Standard benchmark -reports use their own load/timing protocol and do not establish compact routing -by themselves. For a controlled routing-checked comparison, use the optional -helper below. Neither protocol is interchangeable with historical plateau TPS. +Standard benchmark reports use the benchmark's load and timing protocol; they do +not by themselves establish compact routing. The optional helper below runs its +own routing-checked lifetimes with a different timing boundary. See the +[validation reference](CONFIDENCE_VALIDATION.md#validation-scope) for the tested +scope and interpretation of results. -## 4. Generate new artifacts when useful +## 4. Generate artifacts when needed -You can regenerate either artifact independently. Confidence scheduling does not -require temperature fitting before cost profiling or benchmarking. +You can generate costs, temperature calibration, or both. Neither requires +generating the other first. -### Whole-step costs +### Whole-step cost table -Follow [Offline cost profiling](COST_PROFILING.md) and -[Cost export](COST_EXPORT.md). Declare native V0 and legal compact cells at the -chosen K and rank-local graph ladder, enable `enable_iter_perf_stats: true` in -the profiling YAML, and collect with `--cost-scope whole-step`. Reduce the -completed all-rank observations, then export the schema-v2 runtime table with -`dspark_cost_export.py`. - -Use directly measured complete-loop costs; graph-replay medians are diagnostics, -not serving prices. The current whole-step exporter supports attention-DP fanout -or a single rank, not multi-rank attention-DP-off export. The linked guides give -the grid, sample floors, output files and exact commands. +Follow the complete [cost-generation recipe](COST_PROFILING.md): configure a +physical-K-specific native/compact grid, profile complete executor steps, reduce +all-rank observations, and export a schema-v2 runtime table. Graph-replay timing +is a separate diagnostic, not a serving price. The +[cost format/export reference](COST_EXPORT.md) describes structural reuse checks. ### Optional temperature calibration -The paired collector uses a separate tokenized input format: -each JSONL line has a nonempty integer `prompt_token_ids` list. Set the shell -variables below to the inputs chosen in step 1; `RUNTIME_YAML` is the matching -runtime YAML with overlap scheduling enabled. `SOURCE_ROOT` is your checkout, -`RUN_ROOT` is private writable storage, and the timeout reserves the full -startup/model/shutdown budget. +Without calibration, the head uses raw sigmoid confidence probabilities +(unit temperatures). Omit both `confidence_sts_path` in runtime YAML and +`--calibration` in helper commands to select this behavior. Omitting only the +CLI option does not remove a calibration path already present in YAML. + +To generate temperatures, use the paired collector followed by the CPU fitter. +The offline/helper input format is separate from the benchmark format: each +JSONL line contains a nonempty integer `prompt_token_ids` list. + +Set `MODEL_DIR`, `SOURCE_ROOT`, `K`, `RANKS`, `BS` and output work to your chosen +inputs. `BASELINE_YAML` must be a complete matching, confidence-OFF runtime YAML +with no confidence artifact/freshness/mixed/fused fields, explicit topology and +capacity, padded graphs, and overlap scheduling enabled. `RUN_ROOT` is a new +writable directory; `FULL_BUDGET_SECONDS` covers startup, model work and shutdown. ```bash mpirun -n "$RANKS" python tensorrt_llm/tools/dspark_reproduction.py run \ - --mode collect --model "$MODEL_DIR" --config "$RUNTIME_YAML" \ + --mode collect --model "$MODEL_DIR" --config "$BASELINE_YAML" \ --inputs "$INPUTS_JSONL" --source-root "$SOURCE_ROOT" \ --physical-k "$K" --batch-size "$BS" --output-tokens "$OUTPUT_TOKENS" \ --timeout "$FULL_BUDGET_SECONDS" --output "$RUN_ROOT/calibration-observations" @@ -170,49 +171,55 @@ python tensorrt_llm/tools/dspark_calibration.py \ --physical-k "$K" --output "$RUN_ROOT/calibration.json" ``` -The collector runs native full-K verification and joins real confidence logits -to the later accepted-prefix labels for the same proposal and request lifetime. -The fitter requires at least 1000 genuine paired rows by default and produces -per-position temperatures. Collection/fitting is not head training or task -accuracy evaluation. See the [pairing and output contract](CONFIDENCE_VALIDATION.md#paired-collection-and-developer-checks). +Collection runs full-K verification and pairs confidence logits with the later +accepted-prefix labels for the same proposal/request lifetime. The fitter needs +at least 1,000 genuine paired rows by default and produces per-position +temperatures, input hashes and before/after ECE. It is not head training or +task-accuracy evaluation. The [pairing contract](CONFIDENCE_VALIDATION.md#paired-collection-and-developer-checks) +describes the observations. + +To use the result, add `confidence_sts_path: /path/to/calibration.json` under +`speculative_config` in the dynamic runtime YAML, or add +`--calibration "$RUN_ROOT/calibration.json"` to a dynamic helper command. -## 5. Optional controlled comparison and routing checks +## 5. Optional fixed-cohort comparison and routing checks -`dspark_reproduction.py` is for paired collection and controlled fixed-cohort -experiments. It is not a required serving launcher. Its `compare` command -consumes its **own** `rank-N-run.json` records, not `trtllm-bench` reports. -To use it, run its own independent static and dynamic lifetimes with the same -inputs, configuration, output work, warm-up count and budget: +`dspark_reproduction.py` is for paired collection and controlled experiments, +not a required inference launcher. Its `compare` command consumes its own +`rank-N-run.json` records, not `trtllm-bench` reports. + +Use the clean `BASELINE_YAML` and offline input format described above. The helper +requires TP equal to the MPI world size, PP=CP=1 and enough rank-local capacity +for the global cohort. With attention DP, BS must divide evenly across ranks. +Run independent static and dynamic lifetimes with the same inputs, output work, +warm-up count and full budget: ```bash mpirun -n "$RANKS" python tensorrt_llm/tools/dspark_reproduction.py run \ - --mode static --model "$MODEL_DIR" --config "$RUNTIME_YAML" \ + --mode static --model "$MODEL_DIR" --config "$BASELINE_YAML" \ --inputs "$INPUTS_JSONL" --source-root "$SOURCE_ROOT" \ --physical-k "$K" --batch-size "$BS" --output-tokens "$OUTPUT_TOKENS" \ - --timeout "$FULL_BUDGET_SECONDS" --output "$RUN_ROOT/static" + --warmup-cohorts 1 --timeout "$FULL_BUDGET_SECONDS" --output "$RUN_ROOT/static" mpirun -n "$RANKS" python tensorrt_llm/tools/dspark_reproduction.py run \ - --mode dynamic --model "$MODEL_DIR" --config "$RUNTIME_YAML" \ + --mode dynamic --model "$MODEL_DIR" --config "$BASELINE_YAML" \ --inputs "$INPUTS_JSONL" --source-root "$SOURCE_ROOT" \ --physical-k "$K" --batch-size "$BS" --output-tokens "$OUTPUT_TOKENS" \ --cost-table "$COST_TABLE_JSON" \ - --timeout "$FULL_BUDGET_SECONDS" --output "$RUN_ROOT/dynamic" + --warmup-cohorts 1 --timeout "$FULL_BUDGET_SECONDS" --output "$RUN_ROOT/dynamic" python tensorrt_llm/tools/dspark_reproduction.py compare \ --runs "$RUN_ROOT/static/rank-0-run.json" "$RUN_ROOT/dynamic/rank-0-run.json" \ --output "$RUN_ROOT/comparison.json" ``` -Calibration is deliberately omitted above; optionally add -`--calibration "$CALIBRATION_JSON"` to the dynamic command. The helper records -resolved/runtime enablement, native/compact replay, source/artifact identity, -outputs and shutdown. An enabled policy may legally choose native throughout; -`compare` refuses to call that a reproduced compact gain. - +The helper records resolved/runtime enablement, native/compact replay, outputs, +source/artifact identity and shutdown. `compare` requires actual compact routing +for a compact-gain claim; an enabled all-native run is not reported as one. Its TPS covers fixed-cohort enqueue through final response, including prefill -and drain. Use balanced fresh-lifetime repetitions for a stable estimate and -equal static/dynamic record counts for `compare`. Token equality is reported -diagnostically, not substituted for task-accuracy tests. +and drain. Use balanced fresh-lifetime repetitions and equal static/dynamic +record counts. Do not pool these percentages with standard benchmark or plateau +measurements. Token equality is diagnostic, not a task-accuracy test. -See [Validation and developer reference](CONFIDENCE_VALIDATION.md) for the -completed evidence scopes, advanced mixed/freshness contracts and optional tests. +Advanced configuration and optional developer checks are in the +[validation reference](CONFIDENCE_VALIDATION.md). diff --git a/examples/models/core/deepseek_v4/CONFIDENCE_VALIDATION.md b/examples/models/core/deepseek_v4/CONFIDENCE_VALIDATION.md index ef87b5c963f0..affda300eb3a 100644 --- a/examples/models/core/deepseek_v4/CONFIDENCE_VALIDATION.md +++ b/examples/models/core/deepseek_v4/CONFIDENCE_VALIDATION.md @@ -4,10 +4,11 @@ # Confidence scheduling: validation and developer reference Start with the [usage guide](CONFIDENCE_REPRODUCTION.md) for artifact reuse, -self-generation and ordinary inference/benchmarking. This reference separates -validation history and advanced contracts from that user workflow. +self-generation and ordinary inference/benchmarking. This reference records +supported/tested scope and advanced contracts. Developer checks and strict +validation are opt-in, not an ordinary pre-run checklist. -## Recorded validation scope +## Validation scope These are the distinct integration scopes recorded before this documentation reorganization. Original source/native/artifact provenance remains unchanged. @@ -22,6 +23,10 @@ reorganization. Original source/native/artifact provenance remains unchanged. Other prepared batch sizes and configurable physical K are interface coverage, not additional measured results. Existing source tests and historical measurements must not be relabeled as current-head GPU validation. +The interfaces support the global BS128/256/512/1024/1536/2048 grid and smaller +points when actual graph/K/cost/capacity coverage permits. With eight attention-DP +ranks these map to G16/G32/G64/G128/G192/G256; with attention DP off, G is the +global cohort. This mapping is not a GPU-validation claim for every point. The controlled helper reports fixed-cohort enqueue-to-final-response TPS, including prefill and drain. Historical plateau measurements and ordinary @@ -76,10 +81,11 @@ not a blanket DSpark runtime restriction. Replay-entry timing remains a separate diagnostic; it is not renamed to a complete-loop cost or summed with overlapping host/GPU spans. -Optional temperature vectors must have finite positive values for each physical -draft position. Omitting calibration uses raw sigmoid probabilities. Reuse -retains original artifacts and records the actual destination runtime separately; -it does not rewrite old fingerprints or guarantee positive performance. +Temperature vectors, when supplied, must have finite positive values for each +physical draft position. See the [optional calibration workflow](CONFIDENCE_REPRODUCTION.md#optional-temperature-calibration) +for collection, fitting and use. Reuse retains original artifacts and records +the actual destination runtime separately; it does not rewrite old fingerprints +or guarantee positive performance. ### Requested configuration versus execution @@ -224,17 +230,27 @@ source/runtime identity. The fitter retains shard hashes, sample count, temperatures and before/after ECE; degenerate positions are reported rather than promoted as informative fits. This is not task-accuracy qualification. -The optional constructor test below uses the real argument validator and -internal factory seam without constructing a model. Run it only in a compatible -installed native runtime; imports may require a driver. It is a developer -regression, not a colleague pre-run gate or GPU qualification: +Artifact, calibration and observation controls can run with CPU inputs: + +```bash +python -B tests/unittest/tools/test_dspark_cost_export.py +python -B tests/unittest/tools/test_dspark_calibration.py +python -B tests/unittest/tools/test_dspark_observations.py +``` + +The constructor test uses the real argument validator and internal factory seam +without constructing a model. Run it in a compatible installed native runtime; +imports may require a driver: ```bash python -m pytest -q tests/unittest/executor/test_dspark_reproduction_factory.py ``` `executor_cls` is internal and remains rejected as a public LLM argument. -Additional synthetic contract regressions are: +The factory tests cover observation/profile construction, the imported replay +keyword/default contract, position-offset forwarding and routing bookkeeping. +These CPU/native-import controls are not real MPI, GPU execution, accuracy or +performance results. Additional contract regressions are: - `tests/unittest/_torch/speculative/hw_agnostic/test_dspark_workload_admission.py` - `tests/unittest/_torch/speculative/hw_agnostic/test_dspark_admission_resolution.py` diff --git a/examples/models/core/deepseek_v4/COST_EXPORT.md b/examples/models/core/deepseek_v4/COST_EXPORT.md index 50cd42851041..290a0d183fb2 100644 --- a/examples/models/core/deepseek_v4/COST_EXPORT.md +++ b/examples/models/core/deepseek_v4/COST_EXPORT.md @@ -1,137 +1,72 @@ -# Handing measured costs to DSpark - -The offline graph collector and runtime cost consumer have different contracts. -The collector's `verify-cost-measured-summary-v1` output measures the original -graph replay entry, including static input copies. The consumer's schema-v2 -`step_time_ms` represents direct whole-step costs. Do not rename the replay -median to `step_time_ms`, sum potentially overlapping graph/host components, or -construct a new engine fingerprint from the partial loaded-runtime observation. - -The explicit `whole-step` profile mode reuses the executor's existing -`hostStepTimeMS` boundary. It joins the just-finished complete loop to its actual -executor iteration; it does not sum graph and host spans, introduce new CUDA -events, alter streams or add synchronization. The raw exporter retains all-rank -joins, warm-up exclusions and shutdown closure. Graph timing is retained only as -a separately labeled diagnostic in this mode. - -`dspark_cost_export.py` maps those direct whole-step medians into the existing -runtime schema without filling missing cells. It also accepts an existing -schema-v2 artifact and can copy its original bytes after structural validation. -Both paths use the actual runtime schema validator. Plain graph summaries remain -ineligible; that does not prevent reuse of structurally valid supplied costs. - -Direct whole-step collection/export has been exercised in the recorded earlier -physical-K3/G64, global-BS512 environment. Original measurements and fingerprints -are preserved, not relabeled as fresh measurements for later destinations or -sources. See [the reproduction guide](CONFIDENCE_REPRODUCTION.md) for the distinct -validation scopes and remaining limits. - -For a previously measured cost table and a runtime capturing G16 and G64: +# Cost format, export and reuse -```bash -python tensorrt_llm/tools/dspark_cost_export.py \ - --input /path/to/recorded-costs.json \ - --physical-k 3 --graph-batch-sizes 16 64 -``` +`tensorrt_llm/tools/dspark_cost_export.py` accepts either a reduced whole-step +summary or an existing schema-v2 runtime cost table. It validates the requested +physical K and graph ladder using the runtime's structural validator. -`K` is configurable and has no fixed pilot restriction. `G` is the rank-local -request count of the actual captured graph, not always the global batch size. -With eight attention-DP ranks, BS128 and BS512 correspond to G16 and G64; with -attention DP off they do not. Pass the exact graph ladder the engine resolves. -Default reuse permits that configured ladder to be a subset of the originally -measured ladder. The complete original table/fingerprint and native V0, positive -finite timing values, physical K, cell provenance and configured graph coverage -must satisfy the existing consumer. Any drain-guard tail must remain measured -and captured; optional strict freshness still requires the original full ladder. -Missing cells are not interpolated, copied from another K or filled with zero. -No cluster name or site path is used. - -Add `--output /path/to/new-costs.json` only to deliver a new byte-identical copy. -The destination must not exist. No original values, recorded source head or -engine fingerprint are rewritten. In this reuse path there is no exact-head -freshness check and no economic admission requirement; the result warns that -destination freshness and performance are unchecked. Confidence scheduling is -still default-off and must be explicitly enabled through its normal configuration. - -The structural-check report is not a calibration result, GPU health check, -task-accuracy certificate, positive admission receipt or benchmark. Genuine -enabled routing, output checks and a matched static/dynamic comparison remain -separate measurements. The source-only CPU tests do not qualify this complete -clean-checkout procedure on a GPU. - -## Preparing new direct whole-step costs - -Use the configurable K/G/V grid and authorized inputs described in -[COST_PROFILING.md](COST_PROFILING.md). Set `enable_iter_perf_stats: true` in the -ordinary runtime YAML so the existing complete-loop timer is active. Use a -read-only checkout of the original source as `--source-root`; its actual Git -head/diff and loaded module bytes are recorded as provenance, not used as a -destination freshness gate. Runtime/library/GPU observations are preserved in -the raw identity. No new native build is implied by this Python-only tooling. -The Git head is only the original base: the raw `loaded_source_sha256` roster -and recorded runtime snapshot retain executed local Python overlays, including -untracked tooling. A dirty execution is not described as pristine just because -it has the same base head; original tracked diff bytes are hashed separately. - -The additional arguments to the normal profile interface are: +For the complete GPU profile → CPU reduce → CPU export procedure, use +[Generate DSpark whole-step costs](COST_PROFILING.md). For normal runtime +configuration, start with the [reproduction guide](CONFIDENCE_REPRODUCTION.md). -```bash -python -m tensorrt_llm.tools.verify_costs profile \ - --model /path/to/authorized/checkpoint --config /path/to/runtime.yaml \ - --grid /path/to/profile-grid.json --physical-k 3 \ - --cost-scope whole-step --source-root /path/to/original-public-checkout \ - --inputs /path/to/tokenized-prompts.jsonl \ - --output /path/to/new-private-profile-directory -``` +## Accepted timing scopes + +| Input | Treatment | +| --- | --- | +| Whole-step `verify-cost-measured-summary-v1` | Requires `original-executor-loop-hostStepTimeMS` from the explicit whole-step producer. Direct retained medians become schema-v2 `step_time_ms`. | +| Existing schema-v2 runtime table | Structurally validated without changing original bytes, measurements or fingerprint. | +| Graph-replay summary | Rejected as a runtime-price source. Replay-entry medians cannot be relabeled whole-step costs. | + +Whole-step observations use the executor's existing complete-loop timer and +exact all-rank iteration joins. The adapter takes the median of retained +per-iteration rank maxima without summing overlapping graph/host spans. Each +declared native/compact cell requires its own direct measurement, with at least +12 retained samples per generation and 24 pooled per cell. + +The exported fingerprint and source/configuration observations describe the +measurement's original execution. Export does not reobserve the destination or +turn loaded-module hashes into a newly measured live fingerprint. -Use this command under the site's normal MPI/container launcher with the actual -GPU count and full startup/generation/export/shutdown budget. No cluster or -account is embedded in the tooling. This is the configurable interface used by -the recorded earlier integration, not a promise of portable performance or -universal destination qualification. No extra sweep is needed before reusing -valid supplied cost/calibration predictions. +## Check or copy an existing runtime table -Reduce the closed raw observations with the existing `verify_costs.py export` -command. Its defaults retain at least 12 exact joins in each retained generation -and 24 pooled per cell. Then use the DSpark adapter: +For example, to use a table with physical K3 and a runtime capturing G16 and G64: ```bash python tensorrt_llm/tools/dspark_cost_export.py \ - --input /path/to/whole-step-measured-summary.json \ - --physical-k 3 --graph-batch-sizes 16 64 \ - --output /path/to/new-costs.json + --input /path/to/costs.json --physical-k 3 --graph-batch-sizes 16 64 ``` -The adapter exports the retained per-iteration rank maximum's median unchanged; -every schema-v2 measurement references the original summary hash. The recorded -fingerprint is preserved, not refreshed by the CPU exporter. It neither creates -drain constants nor measures confidence-policy overhead, calibration quality or -economic admission. The default table policy margin remains 1%; it is not a -timing measurement. The original summary and raw observations remain available -to inspect scope, source/configuration and raw joins. - -The existing schema-v2 consumer checks global graph counts as rank-local G -times `gpu_count`. This new export currently supports attention-DP fanout or a -single rank; multi-rank attention-DP-off needs a separate schema decision and -is explicitly refused rather than recording incorrect global batch sizes. - -## Graph-only exports stay separate - -The default `graph-replay` collector records -`original-graph-replay-entry-including-static-copies` and -`DIAGNOSTIC_PARTIAL_LOADED_RUNTIME`, with -`production_sps_table_eligible: false`. This command deliberately refuses that -summary even if an eligibility flag is manually changed. Only choosing and -actually collecting the explicit whole-step boundary produces whole-step raw -observations. Changing a label does not extend an existing measured scope. - -GPU-forward values and overlap `iterLatencyMS` are not replacements for clean -complete-loop measurements: the former omits host work, while the latter can -span more than one loop. Optional policy-overhead/economic validation stays -separate. The recorded earlier whole-step integration does not by itself qualify -a later runtime's enabled routing, outputs or matched benchmark. See -[the reproduction guide](CONFIDENCE_REPRODUCTION.md) for the distinct completed -and pending validation scopes. +Pass the actual resolved rank-local graph ladder. Default reuse permits a +configured subset of the measured ladder while validating the entire original +table and fingerprint. Native V0 coverage, offered compact layouts, physical K, +finite positive timings and cell provenance must remain compatible. Any priced +drain tail must be both measured and captured. Missing prices are not +interpolated, copied from another K or filled with zero. + +Add `--output /path/to/new-costs.json` for a byte-identical copy of a valid +schema-v2 input. The destination must not exist. Without `--output`, the command +only prints its structural-check report. + +Structural validation does not require matching source HEAD/node or economic +admission. It warns that destination freshness and performance are unchecked. +Confidence remains default-OFF and needs explicit enablement in runtime YAML. +The separately optional [strict freshness/admission mode](CONFIDENCE_VALIDATION.md#optional-strict-freshness-and-admission) +has additional evidence requirements, including the exact original ladder. + +## New whole-step tables + +When the input is a reduced whole-step summary, `--output` writes a new runtime +table rather than copying summary bytes. Every measurement references the +original summary hash. The default `--minimum-predicted-gain 0.01` sets a policy +margin; it is not a timing measurement. + +The current schema checks global graph counts as rank-local G times GPU count. +New whole-step export therefore supports attention-DP fanout or one rank and +explicitly refuses multi-rank attention-DP-off export. This is an export-schema +limit, not a general DSpark runtime limitation. + +This tool does not fit temperatures, measure policy overhead, certify task +accuracy, grant admission or prove a speedup. Those checks remain separate from +the timing-format conversion. diff --git a/examples/models/core/deepseek_v4/COST_PROFILING.md b/examples/models/core/deepseek_v4/COST_PROFILING.md index 6cd7dccf8462..a4c6d0bb636e 100644 --- a/examples/models/core/deepseek_v4/COST_PROFILING.md +++ b/examples/models/core/deepseek_v4/COST_PROFILING.md @@ -1,100 +1,49 @@ -# Offline DSpark graph-cost profiling - -This interface collects and exports measured graph replay costs. The direct -whole-step producer/export and separate paired collector/fitter have completed -one recorded physical-K3/G64, global-BS512 GPU integration in an earlier -Python/native environment. See [the reproduction guide](CONFIDENCE_REPRODUCTION.md) -for the distinct newer-source validation scope and remaining limits. -The default graph-replay collector retains its narrower scope: it does not itself -generate a serving SPS table, paired calibration observations, an admission -receipt or an enabled throughput result. The explicit complete-loop mode and -CPU schema-v2 adapter are described in [COST_EXPORT.md](COST_EXPORT.md). - -The tools contain no private model, dataset, scratch path, historical receipt or -source-head pin. Use your own authorized local checkpoint and tokenized inputs. -Raw observations include prompt and output tokens; store them in an appropriately -access-controlled output directory. -The Linux collector creates only a new directory with mode0700 and new raw files -with mode0600, using exclusive/non-overwriting creation. It does not change -existing paths, ACLs, process umask or host settings. Existing/symlink targets are -rejected. The exercised Linux collector used new private output paths; CPU -permission controls or one Linux execution do not qualify Windows ACLs or every -destination filesystem. - -## Scope and terminology - -- `K` is the physical draft block produced by the model. Set it explicitly with - `--physical-k`; it must equal the runtime's `speculative_config.max_draft_len`. - Profiling never relabels a K5 execution as K3. -- `G` is the rank-local graph request count. -- `V` is the total verifier token budget in that rank-local graph. `V=0` denotes - the native full-block control, not an empty execution. Positive compact cells - require `K>1` and `G <= V < G*(K+1)`; K1 currently supports only native cells. -- The measured scope is **the original graph replay entry, including static - input copies**. It includes work already inside that graph, including the - embedded drafter and optional confidence head. It is not a target-only timer, - a complete confidence-policy price, or an end-to-end generation timer. - -The offline engine uses the original model loader, attention/backend selection, -graph capture/replay, request scheduler, KV allocator/estimator, sampler, token -updates and shutdown. It prescribes a declared verification cell; it does not -enable or bypass production confidence admission. Ordinary requests outside a -prescribed offline generation are rejected, and no serving RPC endpoint is exposed. -The profiler duplicates its control communicator collectively before worker -construction, retains it through shutdown/evidence closure, then frees it. -Original request broadcasts and runtime communicator/thread-local bindings are -unchanged. Multi-rank profiling requires existing MPI `THREAD_MULTIPLE` support; -the collector checks it and does not modify MPI initialization or thread level. - -Each cell runs a fixed cohort of real requests through prefill, generation and -drain, then repeats that generation. It does not time a single synthetic request -replayed in isolation. Only decode steps with exactly `G` requests on every rank -and the declared captured graph key become samples. Initial prefill, first-draft, -peer-not-ready and shorter drain batches retain the native path and are not timed -as that cell. No forced acceptance is allowed. The default reducer excludes -generation 0 and the first two eligible samples of every retained generation. - -## Prepare a matching runtime and workload - -Use a complete Linux checkout containing this tooling. Build/install TensorRT-LLM -using the repository's [build instructions](../../../../docs/source/installation/build-from-source.md), -or reuse genuinely compatible existing native artifacts through the documented -precompiled/editable-install route. A Python-only tooling change does not by -itself require a native rebuild or new cost/calibration measurements. -Do not combine these Python modules with incompatible native interfaces. Retain -the original native provenance and record changed Python source separately. -The collector records -the actual loaded Python/native-library hashes, package version, GPU class, -topology, graph configuration and observed attention-module classes. A module -roster is not proof of every dispatched kernel or of GPU health. -These raw identity observations are also not a complete weight/JIT/source/build -qualification or a production serving live fingerprint. - -Prepare a runtime YAML accepted by `TorchLlmArgs`, including: - -- `tensor_parallel_size` matching the MPI world size, `pipeline_parallel_size: 1` - and `context_parallel_size: 1`; -- the intended expert parallelism and `enable_attention_dp` setting; -- enough `max_batch_size`, `max_num_tokens`, `max_seq_len` and KV capacity for the - full cohort and requested output length; -- padded decode CUDA graphs whose **resolved rank-local ladder** exactly matches - all `G` values in the profile grid; -- an embedded DSpark configuration with the selected physical K, confidence - scheduling and the optional fused scheduler disabled, no draft-length schedule, - and no confidence SPS/fingerprint/admission/calibration fields. - -This initial driver is aggregate, fixed greedy generation. It rejects LoRA, -guided decoding, custom logits processors, disaggregated KV transfer/connectors, -encode-only execution and in-graph sampling. Prefill graph capture is not part of -this interface; resource-governor control is also unsupported. Keep the intended -numerical/quantization/backend configuration; -do not change it merely to make profiling launch. - -For example, an attention-DP eight-rank profile of global BS512 uses G64. One -possible **declared grid**, not a measured cost table, is: +# Generate DSpark whole-step costs + +Use this recipe when you need a new cost table. Reusing a compatible supplied +table does not require profiling again. Installation, artifact reuse and +ordinary execution are covered by the [reproduction guide](CONFIDENCE_REPRODUCTION.md). + +The workflow is **profile → reduce → export**. Whole-step collection uses the +executor's complete-step `hostStepTimeMS` timer rather than only graph replay. +It uses the normal model, KV, sampler and +shutdown paths. No forced acceptance is permitted. + +## 1. Configure the runtime and grid + +K is the physical draft block, G is the rank-local request count, and V is the +total verifier-row budget of a captured compact layout. V0 is the sentinel for +native full-width `G*(K+1)` verification, not zero work. Every G needs a native +V0 control. Compact cells require K > 1 and `G <= V < G*(K+1)`; choose the compact +cells you want to price rather than sweeping every possible V. + +Prepare a complete runtime YAML for the intended model, precision and backend. +It must have: + +- `tensor_parallel_size` equal to the MPI world size, PP=1 and CP=1; +- the intended expert parallelism and attention DP, plus sufficient request, + token, sequence-length and KV capacity; +- `enable_iter_perf_stats: true` for the complete-loop timer; +- padded CUDA graphs with a resolved rank-local ladder exactly matching all G + values in the grid; +- embedded DSpark at the chosen physical K, with confidence scheduling and the + fused scheduler OFF, no mixed evidence or draft-length schedule, and no + SPS, STS, fingerprint or admission fields. + +Do not use an enabled `dynamic.yaml` unchanged for profiling. The grid's +`record_confidence_head: true` records the loaded head inside the graphs without +enabling serving confidence policy. + +The current whole-step exporter supports attention-DP fanout or a single rank; +multi-rank attention-DP-off export is refused. The offline driver supports +aggregate greedy generation and rejects guided decoding, LoRA, custom logits +processors, disaggregated KV transfer/connectors, in-graph sampling, encode-only +and resource-governor control. Multi-rank operation needs MPI `THREAD_MULTIPLE`. + +For example, use this `profile-grid.json` for K3/G64: ```json { @@ -110,185 +59,114 @@ possible **declared grid**, not a measured cost table, is: } ``` -Use `cuda_graph_config.batch_sizes: [64]` and `enable_padding: true` for that -single-G example; the original runtime's token/request limits must retain G64. -If the resolved ladder has more G values, include their native V0 controls and -desired compact cells. Missing or silently clamped coverage is an error. - -The interface can prepare all previously used global batch sizes -128/256/512/1024/1536/2048. With eight attention-DP ranks these map to -G16/G32/G64/G128/G192/G256. With attention DP off, G is the full global batch, -not BS/8. Smaller points are explicit G inputs too; no pilot batch is hardcoded. -Each additional G requires its own captured graph and measured compatible cells. -Preparation does not imply validation of every batch size. The recorded K3/G64, -global-BS512 integration has completed; other listed points remain prepared -interfaces unless separately identified as tested. Compatible artifact reuse -does not require a new batch-size sweep. - -The JSONL input file must contain at least the largest global cohort size, with -one object per line and a nonempty list of integer `prompt_token_ids`. Use the -actual licensed tokenizer and workload. The driver assigns stable input IDs and -uses the same prompt prefix for repeated cells; prompts and requested output -length must fit `max_seq_len`. - -## Development collection command - -After the complete launch procedure, matching native build and GPU health are -qualified, invoke the CLI under the site's normal MPI/container launcher. For -an already configured, healthy eight-GPU Linux environment the interface is: +This declares shapes, not measured prices. Merge the following fragment into +the complete working runtime YAML for this example; it is not a standalone +configuration: + +```yaml +tensor_parallel_size: 8 +pipeline_parallel_size: 1 +context_parallel_size: 1 +moe_expert_parallel_size: 8 +enable_attention_dp: true +enable_iter_perf_stats: true +cuda_graph_config: + enable_padding: true + batch_sizes: [64] +speculative_config: + decoding_type: DSpark + max_draft_len: 3 + enable_confidence_scheduling: false +``` + +Change K and the grid/graph ladder together for another workload. Each additional +G needs its own native control and desired compact cells. Request/token limits +must retain the requested ladder: a silently clamped graph is not its requested +shape. `max_samples` must accommodate the run's raw replay records. + +The JSONL workload must contain at least the largest global cohort, with a +nonempty integer `prompt_token_ids` list per line. Use the model's tokenizer and +ensure prompt plus output fits `max_seq_len`. Each cell repeats the same real +prompt cohort through prefill, generation and drain. Only decode steps with +exactly G requests on every rank and the declared graph key are samples; +bootstrap and shorter drain steps retain native execution. + +## 2. Profile complete steps on GPUs + +Run from the source checkout using its compatible installed runtime and the +normal allocation/container/MPI procedure. Set these Bash variables to your +paths and selected inputs; for the example, `K=3`, `RANKS=8` and `GRAPH_G=(64)`. +`SOURCE_ROOT` is the checkout used for execution; `PROFILE_YAML` is the complete +confidence-OFF configuration above. Output parents must exist and output paths +must be new. ```bash -mpirun -np 8 --bind-to none python -m tensorrt_llm.tools.verify_costs profile \ - --model /path/to/authorized/checkpoint \ - --config /path/to/runtime.yaml \ - --grid /path/to/profile-grid.json \ - --physical-k 3 \ - --inputs /path/to/tokenized-prompts.jsonl \ - --output /path/to/new-private-output-directory \ - --generations 3 --output-tokens 768 --timeout 12180 +mpirun -n "$RANKS" python -m tensorrt_llm.tools.verify_costs profile \ + --model "$MODEL_DIR" --config "$PROFILE_YAML" --grid "$GRID_JSON" \ + --physical-k "$K" --cost-scope whole-step --source-root "$SOURCE_ROOT" \ + --inputs "$INPUTS_JSONL" --output "$PROFILE_OUT" \ + --generations 3 --output-tokens 768 --timeout "$GENERATION_BUDGET_SECONDS" ``` -These paths are explicit inputs, not supplied artifacts. `--timeout` is the total -generation-phase request budget after construction/identity, not a guarantee that -startup fits a scheduler lease. Reserve the full startup, model, export and -shutdown budget in the outer launcher. Do not restart on quiet logs or an -observation timeout. The output directory must be new and its parent must exist; -no existing output is overwritten. +`--timeout` is the shared generation-phase budget after construction and identity +checks, not the whole allocation budget. Reserve startup, generation, export and +natural shutdown in the outer launcher. Keep the full chosen output workload; +a short run may not meet the sample floors below. -The result contains one `rank-R-record-N.json` per cell/generation/rank. It records -actual replay keys, lease-completed GPU events, request/sampler/KV lifecycle joins, -fixed-length original responses and closure. Completed evidence is retained on -failure with `run_complete: false`; a failure record does not qualify costs. -Successful shutdown is required before raw records can be exported as complete. +The output contains `rank-R-record-N.json` for each cell/generation/rank, with +actual graph keys, per-step samples, request/sampler/KV joins, responses and +source/runtime/model/topology identity. Only completed runs with successful +shutdown can be exported as complete measurements. -## Export and inspect raw measured costs +## 3. Reduce all-rank observations on CPU -Reduction is stdlib-only and can run without importing native TensorRT-LLM: +This reduction uses the standard library and does not import native TensorRT LLM: ```bash python tensorrt_llm/tools/verify_costs.py export \ - --raw /path/to/new-private-output-directory/rank-*-record-*.json \ - --output /path/to/new-measured-summary.json \ + --raw "$PROFILE_OUT"/rank-*-record-*.json --output "$SUMMARY_JSON" \ --exclude-generations 1 --warmup-samples 2 \ --minimum-generation-samples 12 --minimum-samples 24 ``` -The reducer joins exact replay ordinals and executor iterations across **all** -ranks, takes each joined rank maximum, then reports the retained median. It does -not average ranks or sum overlapping spans. It preserves raw-file hashes, -per-rank samples and every exclusion. It rejects duplicate JSON keys, nonfinite -or nonpositive durations, unequal scopes/identities, missing declared cells or -generations, missing/duplicate ranks, reused/reset ordinals, changed observed -layouts and incomplete generation/run/shutdown closure. -Each declared cell must observe one consistent graph key over its generations. -If the workload crosses different sequence-length/attention graph regimes, the -collector refuses to pool them; prepare separate appropriately declared runs. -Each retained generation needs at least 12 exact samples after exclusions, and -each cell needs at least 24 pooled retained samples. A large count in one -generation does not waive another generation's floor. - -Cost keys contain physical K, graph G/V, attention layout, captured-head mode and -the observed graph key. Rows with equal counts but different layouts are not -pooled. Inspect the source/runtime/model/topology identity before comparing or -reusing a summary. Raw summary structure and timing scope remain different from -the runtime's schema-v2 prediction table; renaming a file does not convert it. - -## Fit existing paired confidence observations on CPU - -The standalone fitter accepts tensor-only `.pt` shards containing -`logits: [N, K]`, binary cumulative `prefix_mask: [N, K]` and -`meta.pairing: "draft_seq_ring"`. The labels must be observed against the same -stamped draft sequence and request incarnation as the raw logits, with full -physical-K verification. A shortened compact verification cannot provide labels -for its unverified suffix. This producer contract is not inferred from tensor -dimensions or from a metadata string alone. - -Fit the established 41-point, left-to-right cumulative-survival ECE objective: +The shell expands the raw-file glob. Generation 0 is excluded, as are the first +two eligible samples of each retained generation. Every retained generation +needs at least 12 exact samples, and every cell needs at least 24 pooled samples. +More samples in one generation do not waive another generation's floor. -```bash -python tensorrt_llm/tools/dspark_calibration.py \ - --shards /path/to/private/paired-observations-*.pt \ - --physical-k 3 --output /path/to/new-sts.json -``` +The reducer joins replay ordinals and executor iterations across all ranks, +takes each joined rank maximum, then reports the retained median. It does not +average ranks or sum overlapping spans. It rejects missing/duplicate ranks or +cells, changed layouts, unequal identities, nonfinite/nonpositive times and +incomplete closure. Each cell must retain one consistent observed graph key; +different sequence-length/attention graph regimes need separate declared runs. -The shell expands the shard glob. The default floor is 1,000 paired rows; use -representative data rather than interpreting a small fixture as a calibrated -head. Physical K is explicit, not fixed to the profiling pilot. The output -retains both `sts_temperatures` and `temperatures`, the original shard hashes, -sample count and before/after ECE. Degenerate prefix positions are reported, -not promoted as informative fits. Inputs use PyTorch's tensor-only safe loader; -duplicate shards, wrong pairing/width, nonfinite data and invalid prefix order -are rejected. Existing output files are never overwritten. - -This command has CPU math/CLI coverage, including exact agreement with the -established fitter. The public paired collector and fitter have also completed -the recorded earlier GPU integration. They remain distinct from graph-cost -collection, which does not emit paired shards. No private collector command or -environment variable is a public prerequisite. Neither fitting nor cost -measurement certifies GPU freshness, task accuracy or positive scheduling -economics. Ordinary runtime reuse of a valid supplied STS vector remains -available without fitting it again. - -## Next stages and validation boundary - -A measured summary deliberately sets `production_sps_table_eligible`, -`production_admission_qualified`, `policy_overhead_measured` and -`task_accuracy_qualified` to false. Do not rename it to -`confidence_sps_table_path` or create a positive receipt from graph timings alone. - -Confidence scheduling remains default-off. With explicit enablement, the -default reuse path accepts structurally valid supplied costs and calibration, -with clear warnings when freshness and performance economics are unchecked. -It retains physical-K/graph/layout coverage, finite positive values and ordinary -runtime/acceptance safety. Fresh costs/calibration and positive economic -admission are not universal prerequisites to a fair measured comparison. - -The candidate [self-generation workflow](CONFIDENCE_REPRODUCTION.md) includes -whole-step cost export, correctly paired calibration collection/fitting, and -ordinary static/dynamic execution plus comparison. The graph-replay mode keeps -its narrower recorded scope; it does not price online policy, layout preparation -or end-to-end request execution. -Measure those exposed costs separately when evaluating prediction quality or -scheduling headroom. Do not double-count overlapping spans. - -Users choosing `validate_confidence_artifact_freshness: true` additionally need -an actual destination fingerprint and the optional workload-economic checks, -including measured enabled-path overhead. Native decisions pay that overhead -too. A positive receipt must come from genuine measured/replay evidence, never -from graph timings alone. Neither this collector nor the CPU fitter fabricates -one. The public strict-validation producer remains an explicit pending piece. - -A clean-checkout colleague run must ultimately verify requested versus resolved -enablement, genuine compact routing, completed outputs/task checks and a fair -matched static-versus-enabled benchmark. CPU fake-event tests and smoke imports -are not CUDA replay, full-model accuracy, production admission or performance proof. - -The reusable measurement/export module is -`tensorrt_llm/tools/verify_costs.py`; DSpark shape/runtime adapters are under -`_torch/speculative/dspark_cost_profile.py`, -`_torch/speculative/dspark_cost_profiler.py` and the explicit profile-only engine -and request owner under `_torch/pyexecutor`. Future workload adapters can reuse -timing leases, exact-rank reduction and export while supplying their own execution -shape and layout identity. This does not assume unpublished model-specific details. - -Focused developer controls are: +## 4. Export the runtime table on CPU ```bash -python -B tests/unittest/tools/test_dspark_cost_export.py -python -B tests/unittest/tools/test_dspark_calibration.py -python -B tests/unittest/tools/test_dspark_observations.py -python -m pytest -q tests/unittest/executor/test_dspark_reproduction_factory.py +python tensorrt_llm/tools/dspark_cost_export.py \ + --input "$SUMMARY_JSON" --physical-k "$K" \ + --graph-batch-sizes "${GRAPH_G[@]}" --output "$COST_TABLE_JSON" ``` -The artifact and observation suites use CPU data and finite lifecycle controls; -they do not qualify CUDA events or real distributed execution. The factory suite -needs a matching native import environment but does not construct a model. It -covers the ordinary public argument boundary and the explicit observation/profile -constructor seams, including resource-governor rejection for the profile owner. -It also checks the imported parent replay keyword/default contract and preserves -the same position-offset object in native and measured profile replay branches. -It also covers native-collection and enabled-observation graph parents, including -the same position-offset forwarding and routing bookkeeping. -The file contains fourteen parameter cases; its complete matching-runtime pytest run -is distinct from selected-source CPU boundary controls. -These controls are not full-model GPU, real-MPI, accuracy or performance results. +`GRAPH_G` is a Bash array of the resolved ladder, for example `(16 64)` for two +sizes. The exporter maps directly measured whole-step medians to schema-v2 +`step_time_ms`, retaining their original source/fingerprint and summary hash. +It does not invent missing prices or refresh a historical identity. Use the +new table as `confidence_sps_table_path` in the normal runtime configuration. + +See [Cost format, export and reuse](COST_EXPORT.md) for accepted formats and +structural checks. Costs do not calibrate the head, establish task accuracy or +grant economic admission. + +## Advanced: graph-replay diagnostics + +The same `profile` CLI defaults to `--cost-scope graph-replay`. That mode measures +graph replay including static input copies. Use it for +graph-level investigation, omitting the whole-step-only `--source-root` option. +Its replay medians are not serving prices and are rejected by the schema-v2 cost +exporter. GPU-forward timing and overlap `iterLatencyMS` are not substitutes for +the complete-loop boundary, and overlapping host/GPU spans must not be summed. + +Supported/tested scope and optional developer controls are in the +[validation reference](CONFIDENCE_VALIDATION.md). diff --git a/examples/models/core/deepseek_v4/README.md b/examples/models/core/deepseek_v4/README.md index 9cec1eca0463..bed6e1cca423 100644 --- a/examples/models/core/deepseek_v4/README.md +++ b/examples/models/core/deepseek_v4/README.md @@ -477,13 +477,11 @@ configuration. DSpark can use a trained confidence head to choose a measured verifier-token budget while continuing to propose the full physical draft block. Confidence scheduling is experimental and -off by default. Enabling it requires a compatible cost table; optional temperature calibration -can refine the predictions, but is not required to run. +off by default. Enabling it requires a compatible cost table. Start with [Confidence scheduling and reproduction](CONFIDENCE_REPRODUCTION.md) for ordinary -runtime configuration, artifact reuse, self-generation and benchmarking. The -[validation and developer reference](CONFIDENCE_VALIDATION.md) records exercised configurations, -advanced restrictions, optional freshness checks and developer tests. +runtime YAML and benchmarking, artifact reuse or self-generation, and optional temperature +calibration and controlled comparisons. ## Notes and Troubleshooting