From 951a5a263ab503e412914763d99523074e2caf0e Mon Sep 17 00:00:00 2001 From: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> Date: Fri, 9 Oct 2026 15:03:11 -0700 Subject: [PATCH 1/2] [TRTLLM-14620][feat] Add DSpark ragged CUDA graph runtime Signed-off-by: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> --- .../_torch/pyexecutor/cuda_graph_runner.py | 406 ++++++-- .../engine/runners/decoder/runner.py | 907 +++++++++++++++++- .../_torch/pyexecutor/model_engine.py | 55 ++ tensorrt_llm/_torch/speculative/dspark.py | 13 +- .../speculative/dspark_device_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 | 192 ++++ .../hw_agnostic/test_dspark_device_select.py | 71 ++ .../test_dspark_model_engine_runtime.py | 346 +++++++ .../test_device_window_sampling_layout.py | 87 ++ 13 files changed, 2339 insertions(+), 92 deletions(-) create mode 100644 tensorrt_llm/_torch/speculative/dspark_device_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_device_select.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.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..c354da71612a 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( @@ -173,6 +241,12 @@ def __init__(self, config: CUDAGraphRunnerConfig): self.max_supported_batch_size = config.max_cuda_graph_batch_size self.max_beam_width = config.max_beam_width self.spec_config = config.spec_config + self._dspark_confidence_enabled = bool( + self.spec_config is not None and getattr( + self.spec_config, "enable_confidence_scheduling", False)) + self._dspark_trims_submitted_tokens = bool( + self.spec_config is not None and getattr( + self.spec_config, "enable_confidence_scheduling", False)) self.sparse_config = config.sparse_attention_config self.is_encoder_decoder = config.is_encoder_decoder self.enable_encoder_decoder_mixed_cuda_graph = ( @@ -191,6 +265,16 @@ def __init__(self, config: CUDAGraphRunnerConfig): self._pre_replay_hooks: List[Callable[ [KeyType, int, Dict[str, Any], Dict[str, Any]], None]] = [] self._pre_replay_position_ids: Optional[torch.Tensor] = None + # Low window carried by ordinary ragged padding rows. Zero-real exact + # cells may use a second cached dummy at low+1 to represent V=G*q+r. + self.ragged_pad_verify_len = 0 + # Number of high-token rows in the complete padded graph, including the + # already scheduled attention-DP dummy. Zero is the resting value. + self.ragged_zero_real_high_rows = 0 + # Published by ModelEngine when the ragged fit picks a token bucket; + # None is the resting value (no ragged step in flight). + self.agreed_ragged_bucket: Optional[int] = None + self.adp_shape_agreement: Optional[ADPShapeAgreement] = None self.graphs: Dict[KeyType, torch.cuda.CUDAGraph] = {} self.graph_outputs: Dict[KeyType, @@ -198,6 +282,7 @@ def __init__(self, config: CUDAGraphRunnerConfig): self.graph_metadata: Dict[KeyType, Dict[str, Any]] = {} self.memory_pool = config.cuda_graph_mem_pool self.padding_dummy_requests: Dict[int, LlmRequest] = {} + self.secondary_padding_dummy_requests: Dict[int, LlmRequest] = {} self.dynamic_draft_len_mapping = config.dynamic_draft_len_mapping self.shared_static_tensors: Dict[str, torch.Tensor] = {} @@ -430,7 +515,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 +593,39 @@ 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]: + """Total verified token count when the batch uses ragged windows. + + None unless *every* generation request carries a window; None keeps + the original (uniform) key. Config-gated before touching the batch: + this runs on every graph-key build for every model. + """ + if not self._dspark_trims_submitted_tokens: + return None + # Windows can vanish between the fit and key derivation; a ragged key + # over a window-less batch replays stale token-major row maps (IMA). + if any( + getattr(request, "py_verify_len", None) is None + for request in batch.generation_requests): + return None + # Return the fitted bucket, never a fresh sum over the batch: every + # attention-DP rank must key on the same value, and the fit's + # allgathered arithmetic agrees by construction where a local sum only + # agrees by accident. + return self.agreed_ragged_bucket + + @staticmethod + def _local_draft_len(batch: ScheduledRequests) -> int: + """This rank's draft length, for the attention-DP shape comparison. + + Tolerant where :meth:`get_graph_key` asserts: a malformed batch + reports its max, differs from well-formed peers, and the gate + declines the graph. + """ + return max((len(request.py_draft_tokens) + for request in batch.generation_requests), + default=0) + @staticmethod def _get_mrope_position_delta(request: Any) -> Optional[Any]: mrope_position_delta = getattr(request, "py_mrope_position_delta", None) @@ -568,28 +687,69 @@ 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) + agreement = (self.adp_shape_agreement + if self._dspark_confidence_enabled else None) + # The policy agreement covers shape, not sampling tier. When main's + # in-graph sampler is enabled, retain its per-rank tier agreement. + 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: + local_bucket = self._ragged_verify_bucket(batch) + local_draft_len = self._local_draft_len(batch) + if (not can_run_cuda_graph + or batch_size != agreement.graph_batch_size + or local_bucket != agreement.ragged_bucket + or local_draft_len != 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 (adp_graph_agreement and not self._dspark_confidence_enabled): + # Keep main's sampling-tier exchange unchanged when confidence is + # off; there is no ragged shape to reconcile on this path. 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_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_run_cuda_graph or not all_batch_sizes_equal: + if not all_can_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) + elif (adp_graph_agreement and not reuse_agreement): + # Compare every key component a rank can decide for itself: replay + # freezes all_rank_num_tokens, so a shape mismatch is invisible at + # replay time. draft_len is settable rank-locally (acceptance-rate + # gate), so it is compared rather than assumed. -1 = "not + # applicable" so the payload is plain ints on every rank. + local_bucket = self._ragged_verify_bucket( + batch) if can_run_cuda_graph else None + all_can_graph_batch = self.config.dist.tp_allgather_int64([ + can_run_cuda_graph, + 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() + is_all_gen_only = all(all_can_graph[0] + for all_can_graph in all_can_graph_batch) + all_shapes_equal = all(peer[1:4] == all_can_graph_batch[0][1:4] + for peer in all_can_graph_batch) + + if not is_all_gen_only or not all_shapes_equal: + return None, None, None + + forced_sample_type = (local_sample_type if all( + peer[4] == all_can_graph_batch[0][4] + for peer in all_can_graph_batch) else SampleType.FULL) if not self.enabled or not can_run_cuda_graph: return None, None, None @@ -678,6 +838,12 @@ def get_graph_pool(self): return self.memory_pool 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) token_per_generation = key.draft_len + 1 return (key.num_contexts * key.context_query_len + (key.batch_size * self.max_beam_width - key.num_contexts) * @@ -989,6 +1155,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: @@ -1017,11 +1202,19 @@ def _get_padded_batch(self, batch: ScheduledRequests, 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. + agreement = self.adp_shape_agreement + reuse_agreement = bool( + agreement is not None and agreement.reuse_graph_shape + and agreement.matches(batch, agreement.iteration, padded=False)) + if reuse_agreement: + if not can_run_cuda_graph or not agreement.padding_ready: + raise RuntimeError( + "DSpark cached ADP padding agreement lost a local " + "eligibility invariant") + new_batch_size = agreement.widest_batch_size + + if (self.enabled and self.config.enable_attention_dp + and self.config.mapping.tp_size > 1 and not reuse_agreement): 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] @@ -1041,6 +1234,11 @@ 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) + if (reuse_agreement + and padded_batch_size != agreement.graph_batch_size): + raise RuntimeError( + "DSpark cached ADP padding agreement selected a different " + "graph batch size") if batch_size == padded_batch_size: return 0 @@ -1065,12 +1263,55 @@ def _get_padded_batch(self, batch: ScheduledRequests, key=f"cuda_graph_padding_dummy_fallback_{runtime_draft_len}") return 0 - batch.generation_requests.extend([padding_dummy_request] * padding_size) + # The ragged fit decides padding windows before this generic path + # appends cached dummy requests. A zero-real exact cell may use a low + # and a high object to represent a non-divisible V exactly. + secondary_padding_dummy = None + high_padding_rows = 0 + if self.ragged_zero_real_high_rows: + if (batch_size != 1 + or not batch.generation_requests[0].is_attention_dp_dummy + or self.ragged_zero_real_high_rows > padded_batch_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 + if getattr(self.spec_config, "enable_confidence_scheduling", False): + 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 _get_or_create_padding_dummy( - self, resource_manager: ResourceManager, - runtime_draft_len: int) -> Optional[LlmRequest]: + 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 +1319,13 @@ 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] + if variant not in (0, 1): + raise ValueError( + f"unsupported CUDA padding dummy variant {variant}") + cache = (self.padding_dummy_requests + if variant == 0 else self.secondary_padding_dummy_requests) + 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 +1350,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 +1370,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 +1443,33 @@ 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 low/high 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. + Returns True if either member of the low/high dummy pair 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_requests, + self.secondary_padding_dummy_requests, + ) + ] + 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) + # The cached ADP agreement authenticated this dummy's presence. A + # rebalance may release it between iterations; invalidate the local + # agreement before another batch can suppress the padding collective. + self.adp_shape_agreement = None + 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 +1528,19 @@ def preallocate_padding_dummies(self, logger.info( "Pre-allocated the CUDA graph padding dummy request " f"(draft_len={draft_len}) at warmup.") + if (self._dspark_trims_submitted_tokens and getattr( + self.spec_config, "confidence_sps_table_path", None)): + 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}).") def _add_cross_dummy_request( self, dummy_request: LlmRequest, resource_manager: ResourceManager, @@ -1264,11 +1564,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 @@ -1339,6 +1634,9 @@ def clear(self): self.graph_metadata.clear() self._pre_replay_position_ids = None self.padding_dummy_requests = {} + self.secondary_padding_dummy_requests = {} + self.ragged_zero_real_high_rows = 0 + self.adp_shape_agreement = None del self.memory_pool self.memory_pool = None torch.cuda.empty_cache() diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index f09883c04bcf..9313533c3659 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, Sequence from contextlib import contextmanager -from typing import Any +from typing import Any, List, Optional, Tuple import torch @@ -46,7 +46,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, @@ -70,6 +74,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 ( @@ -404,6 +409,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" ) @@ -431,6 +444,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 # Steady-state generation-only prepare cache (non-speculative overlap # decode). Holds the per-request lists that are invariant while the @@ -691,6 +735,8 @@ 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.secondary_padding_dummy_requests = {} + self.cuda_graph_runner.ragged_zero_real_high_rows = 0 if self.mapping.cp_size > 1: cp_type = self.mapping.cp_config.get("cp_type", None) @@ -791,6 +837,8 @@ def _warmup_scheduled(self, resource_manager: ResourceManager, kv_cache_manager) finally: self.cuda_graph_runner.is_warmup_only = False self.cuda_graph_runner.padding_dummy_requests = {} + self.cuda_graph_runner.secondary_padding_dummy_requests = {} + self.cuda_graph_runner.ragged_zero_real_high_rows = 0 self._run_cuda_graph_warmup(resource_manager) log_mem_snapshot("warmup/after_cuda_graph_capture") # Pre-compile DeepGEMM paged_mqa_logits_metadata for every 32-aligned @@ -806,6 +854,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") @@ -1853,7 +1902,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: @@ -1893,6 +1944,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. @@ -2038,7 +2116,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 @@ -2051,6 +2131,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 @@ -2087,6 +2169,7 @@ def _run_capture_pass( finally: self._force_lora_graph_for_capture = None self.cuda_graph_runner.set_capture_sample_type(None) + self.cuda_graph_runner.agreed_ragged_bucket = None if self._lora.cuda_graph_manager is None: lora_graph_cases = [False] @@ -3006,6 +3089,660 @@ 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_device_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() + 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, @@ -3545,6 +4282,9 @@ def _helix_pack_extend(request, group: int) -> int: # will contain previous batch indices of generation requests previous_batch_indices = [] previous_pos_indices = [] + 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: @@ -3564,6 +4304,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): @@ -3571,10 +4316,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 @@ -3603,6 +4350,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 @@ -3626,30 +4376,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 @@ -3658,7 +4414,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) @@ -3668,7 +4424,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) @@ -3917,6 +4673,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( @@ -3960,18 +4719,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_( @@ -3981,7 +4765,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() ) @@ -4010,10 +4805,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 is_ragged_gen + 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, @@ -4021,7 +4820,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() @@ -4109,7 +4908,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( @@ -4339,6 +5152,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() @@ -4383,6 +5200,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 @@ -4745,6 +5569,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()): @@ -4988,6 +5819,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 b9e44d5ad892..d7a8e926538b 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -830,6 +830,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_device_select.py b/tensorrt_llm/_torch/speculative/dspark_device_select.py new file mode 100644 index 000000000000..d8887e0b20f9 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_device_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. +"""Device-side 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 724b2aa15e36..df3a89376dfc 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -1593,6 +1593,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 da00b2aa01bf..562c2bf87c75 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -88,6 +88,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 ad0cb0482ac4..7985df7c24fc 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -768,6 +768,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] @@ -796,6 +797,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 @@ -820,6 +822,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 @@ -867,6 +870,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] @@ -885,6 +889,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, @@ -925,6 +930,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] @@ -967,6 +973,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..c4d1aed5c46d --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph_runtime.py @@ -0,0 +1,192 @@ +# 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 + +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( + CUDA_GRAPH_DUMMY_REQUEST_ID, + ADPShapeAgreement, + CUDAGraphRunner, + KeyType, + cuda_graph_dummy_request_id, +) + + +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 = CUDAGraphRunner.__new__(CUDAGraphRunner) + 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._dspark_trims_submitted_tokens = True + 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 + runner._dspark_trims_submitted_tokens = False + assert runner._ragged_verify_bucket(batch) 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 diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_select.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_select.py new file mode 100644 index 000000000000..2247b6886f65 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_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_device_select import ( + gather_packed_draft_tokens, + select_windows_device, +) +from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig + + +def test_fresh_confidence_spends_budget_on_the_best_real_row(): + result = select_windows_device( + confidence_logits=torch.tensor( + [[10.0, 10.0, 10.0, 10.0, 10.0], [-10.0, -10.0, -10.0, -10.0, -10.0]] + ), + slot_idx=torch.tensor([0, 1, 0, 0]), + num_real=torch.tensor(2), + budget=torch.tensor(2), + graph_num_tokens=8, + cfg=DSparkScheduleConfig(block_size=5, min_verify_len=1), + pad_len=1, + ) + + assert result.verify_lens.tolist() == [4, 2, 1, 1] + assert result.qo_indptr.tolist() == [0, 4, 6, 7, 8] + assert result.req_idx.tolist() == [0, 0, 0, 0, 1, 1, 2, 3] + assert result.kv_correction.tolist() == [-3, -2, -1, 0, -1, 0, 0, 0] + + +def test_stale_row_fails_open_to_neutral_confidence(): + result = select_windows_device( + confidence_logits=torch.full((2, 5), -10.0), + slot_idx=torch.tensor([0, 1]), + num_real=torch.tensor(2), + budget=torch.tensor(1), + graph_num_tokens=5, + cfg=DSparkScheduleConfig(block_size=5, min_verify_len=1), + stamp=torch.tensor([7, 6]), + expected_stamp=torch.tensor(7), + ) + + assert result.verify_lens.tolist() == [2, 3] + + +def test_gather_packed_draft_tokens_excludes_each_bonus(): + packed = gather_packed_draft_tokens( + next_draft_tokens=torch.tensor([[10, 11, 12], [20, 21, 22]]), + batch_slots=torch.tensor([1, 0]), + verify_lens=torch.tensor([3, 2], dtype=torch.int32), + qo_indptr=torch.tensor([0, 3, 5], dtype=torch.int32), + num_real=2, + total_draft_tokens=3, + ) + + assert packed.tolist() == [20, 21, 10] + + +def test_gather_packed_draft_tokens_rejects_negative_size(): + with pytest.raises(ValueError, match="non-negative"): + gather_packed_draft_tokens( + next_draft_tokens=torch.empty((1, 1)), + batch_slots=torch.tensor([0]), + verify_lens=torch.tensor([1], dtype=torch.int32), + qo_indptr=torch.tensor([0, 1], dtype=torch.int32), + num_real=1, + total_draft_tokens=-1, + ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py new file mode 100644 index 000000000000..b1669101c24c --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py @@ -0,0 +1,346 @@ +# 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 types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine + + +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 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..bcab9063a968 --- /dev/null +++ b/tests/unittest/_torch/speculative/test_device_window_sampling_layout.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import types + +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, None], + ) + SpecMetadata.remap_expanded_sampling_params(meta, torch.tensor([0]), 1) + + +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 3b7fd0deb752b242cb5fc377917a71240362046a Mon Sep 17 00:00:00 2001 From: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> Date: Fri, 9 Oct 2026 19:52:06 -0700 Subject: [PATCH 2/2] [TRTLLM-14620][feat] Integrate DSpark confidence scheduling and mixed verification 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 +++++++++++++ .../_torch/pyexecutor/engine/metadata.py | 5 + .../engine/runners/decoder/runner.py | 315 +++++++-- .../_torch/pyexecutor/model_engine.py | 7 + tensorrt_llm/_torch/pyexecutor/py_executor.py | 628 +++++++++++++++++- tensorrt_llm/_torch/speculative/dspark.py | 62 +- .../speculative/dspark_device_select.py | 527 ++++++++++++++- .../speculative/dspark_mixed_evidence.py | 514 ++++++++++++++ .../_torch/speculative/dspark_planner.py | 82 ++- .../_torch/speculative/dspark_schedule.py | 279 ++++++++ .../_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 | 83 +++ .../hw_agnostic/test_dspark_device_select.py | 30 + .../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_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 +++++- 45 files changed, 7060 insertions(+), 188 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 030b34272450..262406affb23 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 7aedaf8a81e6..2a725292ced4 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -83,12 +83,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]: @@ -3246,6 +3257,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 c354da71612a..3555b4204efc 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -366,6 +366,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/metadata.py b/tensorrt_llm/_torch/pyexecutor/engine/metadata.py index c9ed43c82ef9..48068a897eb0 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,9 @@ def build_attention_metadata( if sparse_attention_config is not None else None ) + ragged_metadata_kwargs = {} + if hasattr(attention_backend.Metadata, "enable_ragged_verification"): + ragged_metadata_kwargs["enable_ragged_verification"] = enable_ragged_verification return attention_backend.Metadata( max_num_requests=max_batch_size, max_num_tokens=max_num_tokens, @@ -82,4 +86,5 @@ def build_attention_metadata( cache_indirection=cache_indirection, num_heads_per_kv=num_heads_per_kv, sparse_metadata_params=sparse_metadata_params, + **ragged_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 9313533c3659..62abdfb38a2d 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -448,26 +448,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_device_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 @@ -855,6 +897,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") @@ -2220,7 +2263,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 @@ -2774,6 +2818,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._dspark_confidence_enabled, ) if isinstance(kv_cache_manager, BaseMambaCacheManager): self.attn_metadata.mamba_chunk_size = getattr( @@ -3007,7 +3052,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 @@ -3016,19 +3063,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. @@ -3089,6 +3157,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. @@ -3451,6 +3603,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. @@ -3470,12 +3657,16 @@ def _apply_device_window_prologue(self, inputs, new_tensors_device) -> bool: """ from tensorrt_llm._torch.speculative.dspark_device_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): @@ -3518,18 +3709,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, @@ -3540,26 +3725,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 ) @@ -3591,10 +3817,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_( @@ -3625,7 +3849,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. @@ -3656,7 +3885,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: @@ -4913,6 +5144,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 ( @@ -5466,6 +5698,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 d7a8e926538b..4b2cb1df265d 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -866,6 +866,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 2a7cdd83d5e3..e277088802f7 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -14,8 +14,8 @@ from contextlib import contextmanager, nullcontext from enum import IntEnum from queue import Queue -from typing import (TYPE_CHECKING, Dict, Iterable, Iterator, List, Literal, - Optional, Tuple, Union) +from typing import (TYPE_CHECKING, Callable, Dict, Iterable, Iterator, List, + Literal, Optional, Tuple, Union) import torch from strenum import StrEnum @@ -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 @@ -129,6 +130,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): @@ -611,6 +756,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() @@ -1961,7 +2112,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) @@ -3774,13 +3925,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.""" @@ -3856,6 +4408,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). @@ -5437,7 +6011,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 @@ -5454,21 +6028,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 @@ -5486,7 +6066,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", @@ -5496,8 +6077,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 @@ -6181,6 +6767,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): @@ -8861,7 +9451,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_device_select.py b/tensorrt_llm/_torch/speculative/dspark_device_select.py index d8887e0b20f9..9eb9190c83e7 100644 --- a/tensorrt_llm/_torch/speculative/dspark_device_select.py +++ b/tensorrt_llm/_torch/speculative/dspark_device_select.py @@ -12,15 +12,17 @@ # 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. -"""Device-side verify-window selection: the in-graph prologue. +"""Device-side verify-window selection for the pre-replay 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. +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_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_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 2626af1ebf23..1c1a1099fd80 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3461,6 +3461,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. @@ -3470,9 +3496,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``. @@ -3514,6 +3539,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( @@ -3529,6 +3637,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 @@ -6741,6 +6894,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.""" @@ -7211,6 +7528,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 9b9e887b5b98..ed49a7d27f32 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1674,6 +1674,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", @@ -1833,6 +1843,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 562c2bf87c75..b3646f2f7ea6 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -189,6 +189,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 f68350399cf5..404c0ca5fbcd 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 c4d1aed5c46d..16cd659ec811 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 from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( @@ -13,6 +14,7 @@ KeyType, cuda_graph_dummy_request_id, ) +from tensorrt_llm._torch.pyexecutor.sampler.sampler_common import SampleType class _Batch: @@ -190,3 +192,84 @@ def test_releasing_padding_dummies_invalidates_the_cached_agreement(): assert runner.release_padding_dummy(SimpleNamespace(), 5) assert freed == [low, high] assert runner.adp_shape_agreement is None + + +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_device_select.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_select.py index 2247b6886f65..0a1dc79aff7a 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_select.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_select.py @@ -6,12 +6,42 @@ import torch from tensorrt_llm._torch.speculative.dspark_device_select import ( + DeviceWindowWorkspace, gather_packed_draft_tokens, select_windows_device, ) from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig +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_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..da592b60072a --- /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_device_select as device_select +import tensorrt_llm._torch.speculative.dspark_schedule as dspark_schedule +from tensorrt_llm._torch.speculative.dspark_device_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_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 1c720ea3a246..119c04694c0a 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 @@ -3094,7 +3094,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( @@ -3139,7 +3139,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 @@ -3154,8 +3154,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( @@ -3270,8 +3270,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( @@ -3757,8 +3757,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 @@ -3778,10 +3778,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): @@ -3798,8 +3798,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): @@ -3825,7 +3825,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 @@ -3896,8 +3896,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. @@ -5242,3 +5241,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())