diff --git a/examples/models/core/deepseek_v4/README.md b/examples/models/core/deepseek_v4/README.md index 5672ef03a73e..6c74b942477d 100644 --- a/examples/models/core/deepseek_v4/README.md +++ b/examples/models/core/deepseek_v4/README.md @@ -36,6 +36,7 @@ for how to build TensorRT LLM from source and start a TRT-LLM Docker container. - [Sparse attention](#sparse-attention) - [KV cache](#kv-cache) - [Quantized checkpoints](#quantized-checkpoints) + - [DSpark confidence scheduling (prototype)](#dspark-confidence-scheduling-prototype) - [Notes and Troubleshooting](#notes-and-troubleshooting) @@ -472,6 +473,73 @@ MXFP4 routed MoE expert weights, TensorRT LLM automatically applies the routed-e configuration. +## DSpark confidence scheduling (prototype) + +DSpark always drafts a block of `K` tokens. With confidence scheduling on, a trained confidence head +scores each drafted position, and in decode-only steps the runtime may verify fewer than all +`G * (K + 1)` rows (`G` = requests per rank, `V` = verified rows). It picks `V` from a **cost table** +of measured step times, and only when the predicted accepted tokens per millisecond beat full +verification by `confidence_min_predicted_gain`. Target verification and the acceptance rule are +unchanged; mixed prefill/decode steps always verify every row. The feature is off by default. + +### Enable it + +```yaml +cuda_graph_config: + enable_padding: true + batch_sizes: [16, 32, 64, 128] # rank-local G; the cost table must cover every entry +speculative_config: + decoding_type: DSpark + max_draft_len: 5 + enable_confidence_scheduling: true + confidence_cost_table_path: /path/to/dspark-costs.json + confidence_calibration_path: /path/to/calibration.json # optional + confidence_min_predicted_gain: 0.04 # default +``` + +A static baseline can reuse the same file with `enable_confidence_scheduling: false`; the other +confidence fields are then ignored. `max_draft_len: 1` has no smaller tier and always runs static. + +### Get a cost table + +Step costs depend on the GPU, parallel layout and kernels, so generate the table on the system +and with the runtime YAML you serve with. One file can hold several `K`; the runtime uses the entry +for `max_draft_len`: + +```bash +python -m tensorrt_llm.tools.dspark_profile_costs \ + --model --config runtime.yaml --k 3 5 --output dspark-costs.json +``` + +The tool loads the model once per `K`. For every captured `G` it runs a short wave of requests +(another one while a budget has too few samples), cycles the full and four compact budgets through +steady decode steps, and records the median whole-step time (draft and verify included). Running it +again adds or replaces the profiled `K` and `G` rows in the same file, so `--k 4` or +`--graph-batch-sizes 512` extends an existing table. `--graph-batch-sizes`, `--tier-fractions` and +`--dataset` override the defaults; see `--help`. + +At startup the runtime checks the table against the engine and fails with the command above when: + +- the file is missing, or has no entry for `max_draft_len`; +- a captured `cuda_graph_config.batch_sizes` entry has no row (extra rows are ignored); +- it was measured on another GPU model, world size, TP/EP size or attention-DP setting. + +A table from another TensorRT LLM version only warns. Tables in the earlier single-`K` format +(`schema_version: 2`) still load. + +### Calibration + +`confidence_calibration_path` points to per-position temperatures (`{"sts_temperatures": [...]}`). +A vector longer than `K` is truncated, so one file fitted at the checkpoint's block size serves +every `K`. Without it the raw sigmoid confidence is used. + +### Limitations + +Confidence scheduling requires the embedded DeepSeek-V4 DSpark drafter, padded CUDA graphs, +PP=1 and CP=1. Guided decoding, LoRA, per-request logits post-processors, per-layer dynamic EPLB, +heuristic sparse top-k, the acceptance-rate gate (`acceptance_rate_window_size` / +`acceptance_rate_threshold`) and non-greedy rejection draft sampling are rejected. + ## 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/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index acd3947a26bb..5252ae5b7f75 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -2111,13 +2111,28 @@ 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.""" - 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: + """Prepare and validate the embedded confidence namespace before loading.""" + last = self.dspark_model.mtp_layers[-1] + head = last.confidence_head + 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 not None and not module_weights: + # A checkpoint without a trained confidence head still serves static + # DSpark; drop the untrained head instead of rejecting the load. + last.confidence_head = head = None + 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_dspark_v41.py b/tensorrt_llm/_torch/models/modeling_dspark_v41.py index 32dd58fbf150..4b030bb6e991 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark_v41.py +++ b/tensorrt_llm/_torch/models/modeling_dspark_v41.py @@ -295,9 +295,6 @@ def load_weights(self, weights: dict[str, torch.Tensor], weight_mapper=None, **k missing = [name for name in required if name not in remapped] if missing: raise ValueError(f"V4.1 DSpark checkpoint is missing draft head weights: {missing}") - confidence_prefix = f"mtp_layers.{self.num_stages - 1}.confidence_head." - if not any(name.startswith(confidence_prefix) for name in remapped): - last.confidence_head = None self._prepare_confidence_head_for_load(remapped) DeepseekV41WeightLoader(self.dspark_model).load_weights(remapped) unexpected, unfed = _v41_load_coverage(self.dspark_model, remapped) diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index e237e87c4d7a..1371a868f97e 100644 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -555,6 +555,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] @@ -569,7 +576,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 77448760f5e9..4b41d67d18fc 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -89,6 +89,8 @@ from .seq_slot_manager import SeqSlotManager GB = 1 << 30 +_DEFAULT_INDEX_MAPPER_RESERVED_SLOTS = 1 +_CONFIDENCE_INDEX_MAPPER_RESERVED_SLOTS = 2 # Page sizes of a standalone drafter's separate KV pool (DFlash, standalone # DSpark). The pool keeps the target's kv_cache_config.tokens_per_block when it @@ -107,6 +109,15 @@ _HELIX_AGG_FREE_GPU_MEMORY_FRACTION = 0.6 +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 ceil_div(a: int, b: int) -> int: return (a + b - 1) // b @@ -3524,6 +3535,9 @@ 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 + if getattr(spec_config, "enable_confidence_scheduling", False): + 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 1f1893025ca2..90301376ac1f 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -1598,8 +1598,29 @@ def __init__(self, config: CUDAGraphRunnerConfig) -> None: # The decoder publishes the fitted bucket before graph lookup. self.agreed_ragged_bucket: Optional[int] = None self.adp_shape_agreement: Optional[ADPShapeAgreement] = None + self._policy_iteration: Optional[int] = None self.secondary_padding_dummy_requests: Dict[int, LlmRequest] = {} + def set_policy_iteration(self, iteration: int) -> None: + """Advance the executor-owned epoch before publishing this step's vote.""" + iteration = int(iteration) + if getattr(self, "_policy_iteration", None) != iteration: + self.adp_shape_agreement = None + self._policy_iteration = iteration + + @contextlib.contextmanager + def pad_batch(self, + scheduled_requests: ScheduledRequests, + resource_manager: ResourceManager, + runtime_draft_len: int = 0): + """Consume the shape vote once, including an exceptional forward exit.""" + try: + with super().pad_batch(scheduled_requests, resource_manager, + runtime_draft_len) as padded_requests: + yield padded_requests + finally: + self.adp_shape_agreement = None + def _ragged_verify_bucket(self, batch: ScheduledRequests) -> Optional[int]: """Return the agreed token bucket only when every row has a window.""" # A window-less batch must not replay stale token-major row maps. @@ -1626,7 +1647,8 @@ def _agree_graph_shape( 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)) + and getattr(self, "_policy_iteration", None) is not None + and agreement.matches(batch, self._policy_iteration, padded=True)) if reuse_agreement: if (not can_run_cuda_graph or batch.batch_size != agreement.graph_batch_size or @@ -1675,8 +1697,9 @@ def _get_num_tokens_for_key(self, key: KeyType) -> int: def _reusable_padding_agreement( self, batch: ScheduledRequests) -> Optional[ADPShapeAgreement]: agreement = self.adp_shape_agreement - if (agreement is not None and agreement.reuse_graph_shape and - agreement.matches(batch, agreement.iteration, padded=False)): + if (agreement is not None and agreement.reuse_graph_shape + and getattr(self, "_policy_iteration", None) is not None and + agreement.matches(batch, self._policy_iteration, padded=False)): return agreement return None @@ -1760,6 +1783,7 @@ def reset_padding_dummies(self) -> None: self.secondary_padding_dummy_requests = {} self.ragged_zero_real_high_rows = 0 self.adp_shape_agreement = None + self._policy_iteration = None def _on_padding_dummy_release(self) -> None: # Releasing a dummy invalidates the agreement that certified it. @@ -1768,7 +1792,7 @@ def _on_padding_dummy_release(self) -> None: def _preallocate_extra_padding_dummy(self, resource_manager: ResourceManager, draft_len: int) -> None: - if not getattr(self.spec_config, "confidence_sps_table_path", None): + if not getattr(self.spec_config, "confidence_cost_table_path", None): return if self._get_or_create_padding_dummy(resource_manager, draft_len, diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py index 61d4952837c0..f12a0801b9f9 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py @@ -56,7 +56,7 @@ from .config import DecoderRunnerConfig from .dummy_requests import DummyRequestFactory from .inputs import DecoderInputPreparer -from .metadata import DecoderMetadata, resolve_mamba_metadata_cls +from .metadata import DecoderMetadata, resolve_mamba_metadata_cls, supports_advanced_draft_sampling from .warmup import WarmupAgreement # Context requests per warmup batch whose breakable prefill graph breaks get @@ -924,7 +924,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 supports_advanced_draft_sampling(self._config.spec_config): + _capture_variant("advanced sampling", force_non_greedy=True) if temperature_only_graphs: # Temperature-only batches replay the vocabulary-sharded Gumbel-max # target sampling (the advanced pass records the stock sampler: its diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py index a9fc0625740e..aec3e60d10f1 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py @@ -417,16 +417,26 @@ def __init__( 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) + table_path = getattr(spec_config, "confidence_cost_table_path", None) if self._dspark_confidence_enabled and table_path: - from tensorrt_llm._torch.speculative.dspark_planner import load_runtime_sps_cost_table + from tensorrt_llm._torch.speculative.dspark_planner import ( + CostTableSetup, + load_runtime_sps_cost_table, + ) + from tensorrt_llm.version import __version__ as trtllm_version - self._dspark_sps_cost_table, _ = 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 + minimum_predicted_gain=float(spec_config.confidence_min_predicted_gain), + live_setup=CostTableSetup( + gpu=torch.cuda.get_device_name(self._buffers.input_ids_cuda.device), + world_size=int(self.mapping.world_size), + tensor_parallel_size=int(self.mapping.tp_size), + moe_expert_parallel_size=int(self.mapping.moe_ep_size), + enable_attention_dp=bool(self._config.enable_attention_dp), + tensorrt_llm_version=trtllm_version, ), ) self._dspark_exact_candidate_cells = self._dspark_sps_cost_table.candidate_cells() diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py index bbe234b984ff..e2041fbd757f 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py @@ -41,6 +41,7 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager + from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig from .inputs import DecoderInputBuffers @@ -372,7 +373,11 @@ def update_spec_metadata( def sync_group_all_greedy_sample( - spec_metadata: SpecMetadata, mapping: Mapping, dist: Distributed | None + spec_metadata: SpecMetadata, + mapping: Mapping, + dist: Distributed | None, + *, + confidence_enabled: bool = False, ) -> None: """All-gather the per-rank greedy flags and store the group AND. @@ -387,13 +392,19 @@ def sync_group_all_greedy_sample( ``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 (mapping.enable_lm_head_tp_in_adp and spec_metadata.use_rejection_sampling): + # Confidence needs the same group-uniform rejection decision even without + # LM-head TP. Feature-off execution retains the existing gate. + if not ( + spec_metadata.use_rejection_sampling + and (mapping.enable_lm_head_tp_in_adp or confidence_enabled) + ): return local_flag = bool(spec_metadata.is_all_greedy_sample) - all_flags = dist.tp_allgather_int64([local_flag])[:, 0] - spec_metadata.group_all_greedy_sample = bool(all_flags.all()) + if mapping.tp_size > 1: + all_flags = 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 @@ -401,6 +412,27 @@ def sync_group_all_greedy_sample( spec_metadata.is_all_greedy_sample = spec_metadata.group_all_greedy_sample +def supports_advanced_draft_sampling(spec_config: DecodingBaseConfig | None) -> bool: + """Whether the engine may capture an advanced draft-sampling variant.""" + return not ( + getattr(spec_config, "enable_confidence_scheduling", False) + and spec_config.use_rejection_sampling + ) + + +def validate_dspark_draft_sampling( + spec_metadata: SpecMetadata, + *, + confidence_enabled: bool, +) -> None: + """Reject unsupported group-uniform sampling before graph lookup.""" + if 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 set_spec_metadata_all_rank_num_tokens( spec_metadata: SpecMetadata, spec_all_rank_num_tokens: list[int], diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index e9f9fb9f8be0..bfec863f56a8 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -84,6 +84,7 @@ get_spec_managers, sync_group_all_greedy_sample, update_spec_metadata, + validate_dspark_draft_sampling, ) from .warmup import DecoderWarmup, WarmupAgreement @@ -635,7 +636,16 @@ def _forward_decoder( # unpopulated (greedy) buffers, hanging the run (e.g. MTP nextn>=2). if spec_metadata is not None: spec_metadata.update_is_all_greedy_sample(padded_graph_requests.all_requests()) - sync_group_all_greedy_sample(spec_metadata, self._mapping, self._dist) + sync_group_all_greedy_sample( + spec_metadata, + self._mapping, + self._dist, + confidence_enabled=self._dspark_confidence_enabled, + ) + validate_dspark_draft_sampling( + spec_metadata, + confidence_enabled=self._dspark_confidence_enabled, + ) selection = self._graphs.select( padded_graph_requests, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index ab97eab14825..db0e85063b8e 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -83,6 +83,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 @@ -143,6 +144,137 @@ _UNBOUNDED_STATS_MAX_LEN = -1 +# Per-rank int64 words of the DSpark attention-DP agreement, all-gathered once +# per generation step: +# 0 an exact decision was prepared on this rank (0/1) +# 1 real generation rows on this rank +# 2 this rank's batch can run a CUDA graph (0/1) +# 3 number of measured (G, V) cells +# 4..11 cost-table identity words +# 12 native expected yield, scaled +# 13.. compact expected yield per cell, scaled (-1 = infeasible) +# followed by the trailer: batch size, padding dummy ready, secondary padding +# dummy ready. +_DSPARK_EXACT_WIRE_PREFIX_LEN = 13 +_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]]: + """Decode aggregate native and per-cell compact yields across ranks. + + A cell infeasible on any rank decodes to zero yield, so it never wins. + """ + 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 = {} + 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 + ] + if any(value < 0.0 for value in rank_yields): + compact_expected_yields[verifier_budget] = 0.0 + continue + if any(value > float(payload[_DSPARK_EXACT_NATIVE_YIELD_INDEX]) + for payload, value in zip(payloads, rank_yields)): + raise RuntimeError( + "DSpark exact-policy compact yield exceeds native yield") + compact_expected_yields[verifier_budget] = sum( + rank_yields) / yield_scale + return native_expected_yield, compact_expected_yields + + +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 _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): # MAX reduction gives context precedence when ADP ranks have mixed work. @@ -838,6 +970,12 @@ def _initialize_execution(self, resources: PyExecutorResources) -> None: if window and threshold is not None: self.speculation_gate = SpeculationGate(window, threshold) + self._dspark_cost_profiler = None + if getattr(self.model_engine, "_dspark_confidence_enabled", False): + from ..speculative.dspark_profile import DSparkCostProfiler + self._dspark_cost_profiler = DSparkCostProfiler.from_env( + rank=getattr(self.dist, "rank", 0) if self.dist else 0) + # response used data self.response_lock = threading.Lock() self.response_cv = threading.Condition(self.response_lock) @@ -4424,14 +4562,351 @@ 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) -> None: + buffer = worker.staged_confidence_buffer() + if buffer is None: + return + worker.bump_draft_seq() + 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 + + 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, 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 + # Steps with context requests keep full verification. + if (compute_windows and not scheduled_batch.context_requests + and 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) + + exact_yield_scale = 1_000_000 + 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, + ) + runner = self.model_engine.cuda_graph_runner + local_payload = [ + 1 if exact_local is not None else 0, + int(exact_local.num_requests) if exact_local is not None else 0, + 1 if scheduled_batch.can_run_cuda_graph else 0, + len(exact_cells), + *exact_identity_words, + native_yield_wire, + *compact_yields_wire, + 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, + ] + 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 + if any(len(payload) != expected_payload_len for payload in payloads): + raise RuntimeError( + "DSpark ADP agreement payload has an inconsistent shape") + + exact_cell_counts = [int(payload[3]) for payload in payloads] + exact_identities = [ + tuple(int(value) for value in payload[4:12]) 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[2]) for payload in payloads) + peer_stats = None + exact_policy_agreed = bool(all(int(payload[0]) for payload in payloads)) + if not exact_policy_agreed: + exact_local = None + + 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)) + 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 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: + 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 + + cost_profiler = getattr(self, "_dspark_cost_profiler", None) + if cost_profiler is not None: + # Cost profiling: a fixed schedule replaces the policy. + selected_verifier_budget = cost_profiler.choose( + iteration=int(self.iter_counter), + graph_batch_size=common_graph_batch_size, + budgets=(0, ) + + exact_table.production_candidate_budgets( + common_graph_batch_size), + all_ranks_full=all( + int(payload[1]) == common_graph_batch_size + for payload in payloads), + ) + else: + ( + native_expected_yield, + compact_expected_yields, + ) = _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, + 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) + return int(planner.max_verify_len) + def _handle_dynamic_draft_len(self, scheduled_batch: ScheduledRequests) -> None: - """Synchronize draft length and buffers before preparing resources.""" - self.runtime_draft_len = 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. + self.runtime_draft_len = update_draft_len( + self.model_engine, + scheduled_batch, + speculation_permanently_disabled=self. + speculation_permanently_disabled) + return + + self.model_engine.cuda_graph_runner.set_policy_iteration( + self.iter_counter) + 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.runtime_draft_len = runtime_draft_len + else: + self.runtime_draft_len = 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.""" @@ -4507,6 +4982,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). @@ -6456,7 +6953,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 @@ -6473,21 +6970,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 = self.model_engine.cuda_graph_runner + 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 @@ -6502,10 +7005,11 @@ def _resume_padding_dummies_after_rebalance( that every manager the dummy was registered with is covered, not just the main KV cache manager. """ - runner = self.model_engine.cuda_graph_runner + 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", @@ -6515,8 +7019,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 @@ -7160,6 +7669,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) retention = request.kv_cache_retention_config if self._is_kv_manager_v2 and retention is not None: diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 757b84dc3124..cfb4658168e0 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 @@ -40,6 +41,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). @@ -63,7 +65,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]: @@ -75,6 +80,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. @@ -454,7 +488,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 @@ -658,41 +700,47 @@ def _build_verify_planner(self, draft_model, block_size: int) -> None: cfg = DSparkScheduleConfig(block_size=block_size, min_verify_len=1) - table_path = self.spec_config.confidence_sps_table_path - # ModelEngine owns the actual CUDA graph ladder, so it is the only - # component that may read, authenticate, and construct the deployment - # cost object. Reopening this path here creates a TOCTOU window in - # which graph capture and scheduling can observe different files. - logger.info( - f"DSpark: deferring exact SPS table {table_path} to the validated " - "ModelEngine runtime object" - ) + # ModelEngine loads the cost table because it owns the captured CUDA + # graph ladder; the planner receives that runtime object before its + # first decision. # Calibration lives on the draft's confidence head; fall back to a plain # sigmoid when the head is absent (no confidence to calibrate anyway). head = resolve_confidence_head(draft_model) + if head is None: + logger.warning( + "DSpark confidence scheduling is enabled, but the checkpoint has no " + "confidence-head weights; every step verifies the full draft block." + ) apply_calibration = head.apply_sts if head is not None else None - sts_path = self.spec_config.confidence_sts_path - if sts_path and head is None: + calibration_path = self.spec_config.confidence_calibration_path + if calibration_path and head is None: # Refuse rather than degrade: a configured calibration that cannot # be attached is a wiring bug, never a preference. raise ValueError( - f"confidence_sts_path is set but no confidence head was found " + f"confidence_calibration_path is set but no confidence head was found " f"on {type(draft_model).__name__}; calibration cannot be " f"applied, so the planner would schedule on raw sigmoid " f"survivals. This is a model-wiring bug, not a config choice." ) - if sts_path and head is not None: + if calibration_path and head is not None: # Accepts either key spelling ("sts_temperatures" here, # "temperatures" in SGLang); the vectors are interchangeable. - temps = load_sts_temperatures_from_path(sts_path) - head.load_sts_temperatures(torch.tensor(temps, dtype=torch.float32)) - logger.info(f"DSpark: loaded STS calibration from {sts_path}") + temps = load_sts_temperatures_from_path(calibration_path) + num_positions = int(head.sts_temperatures.numel()) + if len(temps) < num_positions: + raise ValueError( + f"{calibration_path} has {len(temps)} temperatures but the " + f"confidence head scores {num_positions} draft positions" + ) + # Temperatures are per draft position, so a calibration fitted at a + # larger K serves every smaller K through its leading entries. + head.load_sts_temperatures(torch.tensor(temps[:num_positions], dtype=torch.float32)) + logger.info(f"DSpark: loaded confidence calibration from {calibration_path}") self.verify_planner = DSparkVerifyPlanner( cfg=cfg, apply_calibration=apply_calibration, - device_windows=bool(self.spec_config.enable_fused_confidence_scheduler), ) logger.info( f"DSpark verify planner: max_verify_len={self.verify_planner.max_verify_len}, " diff --git a/tensorrt_llm/_torch/speculative/dspark_planner.py b/tensorrt_llm/_torch/speculative/dspark_planner.py index 812a6932b2db..038b5a04627a 100644 --- a/tensorrt_llm/_torch/speculative/dspark_planner.py +++ b/tensorrt_llm/_torch/speculative/dspark_planner.py @@ -12,61 +12,47 @@ # 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. -"""Authenticated DSpark verification costs and exact ``(G, V)`` selection.""" +"""DSpark verification cost tables and exact ``(G, V)`` selection. + +A cost table holds measured whole-step decode times per physical draft length +``K``, rank-local CUDA-graph batch size ``G`` and verifier-token budget ``V``. +One file can hold several ``K`` values; the runtime loads the entry matching +``max_draft_len``. ``tensorrt_llm.tools.dspark_profile_costs`` generates it. +""" import hashlib import json import math -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from pathlib import Path import numpy as np +from tensorrt_llm.logger import logger + _MAX_EXACT_COMPACT_CELLS_PER_G = 4 _MAX_EXACT_COMPACT_CELLS_TOTAL = 32 +COST_TABLE_SCHEMA_VERSION = 3 +PROFILE_COMMAND_HINT = ( + "python -m tensorrt_llm.tools.dspark_profile_costs --model " + "--config --k {k} --output {path}" +) + __all__ = [ - "ConfidenceWorkloadAdmission", - "EXACT_SPS_SELECTOR_IDENTITY_SHA256", - "ExactSpsDrainGuard", + "COST_TABLE_SCHEMA_VERSION", + "CostTableSetup", "ExactSpsCostRow", "ExactSpsCostTable", - "evaluate_confidence_workload_admission", - "load_confidence_workload_admission", + "build_cost_table_payload", "load_runtime_sps_cost_table", + "merge_cost_table_payloads", + "read_cost_table_rows", "select_exact_sps_candidate", - "validate_sps_cost_table_payload", ] -# This digest identifies the marginal selector semantics. Fixed -# confidence-path overhead is deliberately absent: it is a workload-level -# feature-activation cost, evaluated once before engine construction. -_EXACT_SPS_SELECTOR_IDENTITY = ( - "dspark-exact-sps-marginal-selector-v1:immediate-goodput+group-drain:no-fixed-overhead" -) -EXACT_SPS_SELECTOR_IDENTITY_SHA256 = hashlib.sha256( - _EXACT_SPS_SELECTOR_IDENTITY.encode() -).hexdigest() - - -@dataclass(frozen=True) -class ConfidenceWorkloadAdmission: - """Authenticated workload-level decision made before runtime startup.""" - - admitted: bool - physical_k: int - policy_steps: int - compact_choices: int - gross_compact_value_ms_lower_bound: float - fixed_confidence_path_overhead_ms_upper_bound: float - safety_margin_ms: float - net_value_ms_lower_bound: float - admission_identity_sha256: str - receipt_sha256: str - - @dataclass(frozen=True) class ExactSpsCostRow: """Direct whole-step measurements for one rank-local graph size ``G``.""" @@ -87,8 +73,7 @@ def __post_init__(self) -> None: for value in self.token_counts ) step_time_ms = tuple( - _require_positive_finite_number(value, field="SPS step time") - for value in self.step_time_ms + _require_positive_finite_number(value, field="step time") for value in self.step_time_ms ) if any(b <= a for a, b in zip(token_counts, token_counts[1:])): raise ValueError("token_counts must be strictly increasing") @@ -96,67 +81,6 @@ def __post_init__(self) -> None: object.__setattr__(self, "step_time_ms", step_time_ms) -@dataclass(frozen=True) -class ExactSpsDrainGuard: - """Measured policy metadata for conservative group-E2E admission. - - ``mean_output_tokens_per_request_iteration`` must come from a matched - workload trace, while ``tail_graph_batch_size`` identifies the measured - native ``T(G, 0)`` tail cell in the same exact cost table. Keeping both - values in the authenticated table metadata prevents serving code from - silently reusing trace-specific constants on another workload. - """ - - loss_multiplier: float - mean_output_tokens_per_request_iteration: float - minimum_group_value_ms: float - tail_graph_batch_size: int - source_result_sha256: str - - def __post_init__(self) -> None: - loss_multiplier = _require_positive_finite_number( - self.loss_multiplier, field="iteration drain loss_multiplier" - ) - mean_output_tokens = _require_positive_finite_number( - self.mean_output_tokens_per_request_iteration, - field="iteration drain mean_output_tokens_per_request_iteration", - ) - minimum_group_value_ms = _require_nonnegative_finite_number( - self.minimum_group_value_ms, - field="iteration drain minimum_group_value_ms", - ) - tail_graph_batch_size = _require_exact_int( - self.tail_graph_batch_size, - field="iteration drain tail_graph_batch_size", - minimum=1, - ) - source_result_sha256 = _require_sha256( - self.source_result_sha256, - field="iteration drain source_result_sha256", - ) - object.__setattr__(self, "loss_multiplier", loss_multiplier) - object.__setattr__( - self, - "mean_output_tokens_per_request_iteration", - mean_output_tokens, - ) - object.__setattr__(self, "minimum_group_value_ms", minimum_group_value_ms) - object.__setattr__(self, "tail_graph_batch_size", tail_graph_batch_size) - object.__setattr__(self, "source_result_sha256", source_result_sha256) - - def identity_payload(self) -> dict[str, object]: - """Canonical fields included in all-rank table agreement.""" - return { - "loss_multiplier": self.loss_multiplier, - "mean_output_tokens_per_request_iteration": ( - self.mean_output_tokens_per_request_iteration - ), - "minimum_group_value_ms": self.minimum_group_value_ms, - "source_result_sha256": self.source_result_sha256, - "tail_graph_batch_size": self.tail_graph_batch_size, - } - - @dataclass(frozen=True) class ExactSpsCostTable: """Directly measured whole-step costs keyed by exact ``(G, V)``. @@ -173,8 +97,7 @@ class ExactSpsCostTable: tables: dict[int, ExactSpsCostRow] max_draft_len: int - minimum_predicted_gain: float = 0.01 - iteration_drain_guard: ExactSpsDrainGuard | None = None + minimum_predicted_gain: float = 0.04 identity_sha256: str = field(init=False) def __post_init__(self) -> None: @@ -192,14 +115,11 @@ def __post_init__(self) -> None: if any(not isinstance(table, ExactSpsCostRow) for table in normalized.values()): raise TypeError("ExactSpsCostTable values must be ExactSpsCostRow instances") for graph_batch_size, table in normalized.items(): - budgets = tuple( - _require_exact_int(value, field="measured verifier budget", minimum=0) - for value in table.token_counts - ) + budgets = table.token_counts if budgets[0] != 0: raise ValueError( - "Every exact SPS G requires V=0 native static physical-K " - f"fallback; missing for G={graph_batch_size}" + "Every G requires a V=0 native static physical-K measurement; " + f"missing for G={graph_batch_size}" ) max_verifier_budget = graph_batch_size * (max_draft_len + 1) invalid_positive_budgets = [ @@ -209,8 +129,7 @@ def __post_init__(self) -> None: ] if invalid_positive_budgets: raise ValueError( - "Exact SPS positive verifier budgets must satisfy " - "G <= V <= G*(K+1); " + "Positive verifier budgets must satisfy G <= V <= G*(K+1); " f"G={graph_batch_size}, K={max_draft_len}, " f"invalid V={invalid_positive_budgets}" ) @@ -221,7 +140,7 @@ def __post_init__(self) -> None: ] if len(production_budgets) > _MAX_EXACT_COMPACT_CELLS_PER_G: raise ValueError( - "Exact SPS graph budget exceeds the production limit of " + "Cost table exceeds the limit of " f"{_MAX_EXACT_COMPACT_CELLS_PER_G} compact V cells per G; " f"G={graph_batch_size} has {len(production_budgets)}" ) @@ -237,9 +156,9 @@ def __post_init__(self) -> None: ) if production_cell_count > _MAX_EXACT_COMPACT_CELLS_TOTAL: raise ValueError( - "Exact SPS graph budget exceeds the production limit of " - f"{_MAX_EXACT_COMPACT_CELLS_TOTAL} compact (G,V) cells; " - f"artifact has {production_cell_count}" + "Cost table exceeds the limit of " + f"{_MAX_EXACT_COMPACT_CELLS_TOTAL} compact (G,V) cells across captured " + f"graph sizes; it has {production_cell_count}" ) object.__setattr__(self, "tables", normalized) object.__setattr__(self, "max_draft_len", max_draft_len) @@ -247,28 +166,14 @@ def __post_init__(self) -> None: self.minimum_predicted_gain, field="minimum_predicted_gain" ) object.__setattr__(self, "minimum_predicted_gain", minimum_predicted_gain) - iteration_drain_guard = self.iteration_drain_guard - if iteration_drain_guard is not None: - if not isinstance(iteration_drain_guard, ExactSpsDrainGuard): - raise TypeError("iteration_drain_guard must be an ExactSpsDrainGuard or None") - if iteration_drain_guard.tail_graph_batch_size not in normalized: - raise ValueError( - "Iteration drain guard requires a measured native tail table for " - f"G={iteration_drain_guard.tail_graph_batch_size}" - ) # This is the rank-agreement identity used on every exact scheduling # step. It covers the canonical ordered grid, every measured cost, K, - # and the minimum-gain policy. Two artifacts that merely have the same + # and the minimum-gain policy. Two tables that merely have the same # number of cells must not be allowed to index each other's yields. identity_payload = { - "schema": "exact-sps-runtime-v1", + "schema": "exact-sps-runtime-v2", "max_draft_len": max_draft_len, "minimum_predicted_gain": minimum_predicted_gain, - "iteration_drain_guard": ( - iteration_drain_guard.identity_payload() - if iteration_drain_guard is not None - else None - ), "tables": [ { "graph_batch_size": graph_batch_size, @@ -295,7 +200,7 @@ def for_graph_batch_size(self, num_requests: int) -> ExactSpsCostRow: return self.tables[graph_batch_size] except KeyError as error: raise ValueError( - f"SPS cost artifact has no direct measurements for G={graph_batch_size}" + f"Cost table has no direct measurements for G={graph_batch_size}" ) from error def step_times(self, num_tokens: np.ndarray, num_requests: int) -> np.ndarray: @@ -315,7 +220,7 @@ def step_times(self, num_tokens: np.ndarray, num_requests: int) -> np.ndarray: missing = sorted(set(requested_budgets) - measured.keys()) if missing: raise ValueError( - f"SPS cost artifact has no direct measurements for G={num_requests}, V={missing}" + f"Cost table has no direct measurements for G={num_requests}, V={missing}" ) return np.asarray( [measured[value] for value in requested_budgets], @@ -330,12 +235,9 @@ def candidate_budgets( self, num_requests: int, *, include_native: bool = False ) -> tuple[int, ...]: """Return exact per-G V candidates, excluding native by default.""" - budgets = tuple( - _require_exact_int(value, field="measured verifier budget", minimum=0) - for value in self.for_graph_batch_size(num_requests).token_counts - ) + budgets = self.for_graph_batch_size(num_requests).token_counts if include_native: - return budgets + return tuple(budgets) return tuple(value for value in budgets if value != 0) def candidate_cells(self) -> tuple[tuple[int, int], ...]: @@ -357,278 +259,316 @@ def production_candidate_budgets(self, num_requests: int) -> tuple[int, ...]: ) -def _read_sps_cost_payload(path: str | Path) -> dict[str, object]: - with Path(path).open(encoding="utf-8") as file: +@dataclass(frozen=True) +class CostTableSetup: + """Hardware and parallel layout a cost table was measured on. + + Step costs from another GPU model or parallel layout do not describe the + running engine, so the loader rejects them. A different TensorRT-LLM + version only warns: kernels drift by a few percent, but the table stays a + reasonable price list until it is regenerated. + """ + + gpu: str + world_size: int + tensor_parallel_size: int + moe_expert_parallel_size: int + enable_attention_dp: bool + tensorrt_llm_version: str = "" + + def to_json(self) -> dict[str, object]: + return { + "gpu": self.gpu, + "world_size": self.world_size, + "tensor_parallel_size": self.tensor_parallel_size, + "moe_expert_parallel_size": self.moe_expert_parallel_size, + "enable_attention_dp": self.enable_attention_dp, + "tensorrt_llm_version": self.tensorrt_llm_version, + } + + +_SETUP_ERROR_FIELDS = ( + "gpu", + "world_size", + "tensor_parallel_size", + "moe_expert_parallel_size", + "enable_attention_dp", +) + + +def _read_cost_table_file(path: str | Path) -> dict[str, object]: + path = Path(path) + if not path.is_file(): + raise FileNotFoundError( + f"DSpark cost table {path} does not exist. Generate one with: " + + PROFILE_COMMAND_HINT.format(k="", path=path) + ) + with path.open(encoding="utf-8") as file: payload = json.load(file) if not isinstance(payload, dict): - raise TypeError("SPS cost artifact must contain a JSON object") + raise TypeError(f"DSpark cost table {path} must contain a JSON object") return payload -def load_runtime_sps_cost_table( - path: str | Path, - *, - graph_batch_sizes: Sequence[int], - max_draft_len: int, - live_engine_fingerprint_path: str | Path | None = None, -) -> tuple[ExactSpsCostTable, dict[str, object]]: - """Load an authenticated schema-v2 exact ``T(G,V)`` table. +def _parse_rows(rows: object, *, where: str) -> dict[int, ExactSpsCostRow]: + if not isinstance(rows, dict) or not rows: + raise TypeError(f"{where} must be a non-empty object keyed by graph batch size") + parsed: dict[int, ExactSpsCostRow] = {} + for key, row in rows.items(): + graph_batch_size = _parse_graph_batch_size_key(key) + if not isinstance(row, dict): + raise TypeError(f"{where} G={key} must be an object") + missing = sorted({"token_counts", "step_time_ms"} - row.keys()) + if missing: + raise ValueError(f"{where} G={key} is missing {', '.join(missing)}") + token_counts = row["token_counts"] + step_time_ms = row["step_time_ms"] + if not isinstance(token_counts, list) or not isinstance(step_time_ms, list): + raise TypeError(f"{where} G={key} token_counts and step_time_ms must be lists") + parsed[graph_batch_size] = ExactSpsCostRow( + token_counts=tuple( + _require_json_int(value, field=f"{where} G={key} verifier budget", minimum=0) + for value in token_counts + ), + step_time_ms=tuple( + _require_positive_finite_number(value, field=f"{where} G={key} step time") + for value in step_time_ms + ), + ) + return parsed + + +def _legacy_v2_setup(fingerprint: object) -> dict[str, object]: + """Map a schema-v2 engine fingerprint onto the schema-v3 setup fields.""" + if not isinstance(fingerprint, dict): + return {} + setup: dict[str, object] = {} + if isinstance(fingerprint.get("gpu"), str): + setup["gpu"] = fingerprint["gpu"] + if type(fingerprint.get("gpu_count")) is int: + setup["world_size"] = fingerprint["gpu_count"] + topology = fingerprint.get("topology") + if isinstance(topology, str): + try: + topology = json.loads(topology) + except ValueError: + topology = None + if isinstance(topology, dict): + for source, target in ( + ("tp", "tensor_parallel_size"), + ("ep", "moe_expert_parallel_size"), + ("attention_dp", "enable_attention_dp"), + ): + if source in topology: + setup[target] = topology[source] + return setup + - 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. +def read_cost_table_rows( + payload: Mapping[str, object], *, max_draft_len: int, path: str | Path = "" +) -> tuple[dict[int, ExactSpsCostRow], dict[str, object]]: + """Return ``({G: row}, recorded setup)`` for one physical K. + + Schema 3 keys costs by K, then G. Schema 2 (one K per file, written by + the earlier offline exporter) is still read so existing tables keep + working; its provenance digests are not re-verified. """ - payload = _read_sps_cost_payload(path) + k = _require_exact_int(max_draft_len, field="max_draft_len", minimum=1) 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: - raise ValueError( - "Schema-v2 exact SPS artifacts require an independently generated " - "live engine fingerprint path" + if schema_version == COST_TABLE_SCHEMA_VERSION: + costs = payload.get("costs") + if not isinstance(costs, dict) or not costs: + raise TypeError(f"DSpark cost table {path} has no 'costs' object") + if str(k) not in costs: + available = sorted(int(key) for key in costs if str(key).isdecimal()) + raise ValueError( + f"DSpark cost table {path} has no costs for max_draft_len={k} " + f"(it has K={available}). Add them with: " + + PROFILE_COMMAND_HINT.format(k=k, path=path) + ) + setup = payload.get("setup") + return ( + _parse_rows(costs[str(k)], where=f"{path} K={k}"), + dict(setup) if isinstance(setup, dict) else {}, + ) + if schema_version == 2: + fingerprint = payload.get("engine_fingerprint") + recorded_k = fingerprint.get("max_draft_len") if isinstance(fingerprint, dict) else None + if recorded_k != k: + raise ValueError( + f"DSpark cost table {path} was measured for max_draft_len={recorded_k}, " + f"but the runtime uses max_draft_len={k}. Generate a table for K={k} with: " + + PROFILE_COMMAND_HINT.format(k=k, path=path) + ) + return ( + _parse_rows(payload.get("cost_tables"), where=f"{path} K={k}"), + _legacy_v2_setup(fingerprint), ) - live_engine_fingerprint = _read_sps_cost_payload(live_engine_fingerprint_path) - validate_sps_cost_table_payload( - payload, - graph_batch_sizes=graph_batch_sizes, - max_draft_len=max_draft_len, - live_engine_fingerprint=live_engine_fingerprint, + raise ValueError( + f"DSpark cost table {path} has unsupported schema_version={schema_version!r}; " + f"expected {COST_TABLE_SCHEMA_VERSION} (or legacy 2)" ) - return _build_exact_sps_cost_table(payload, max_draft_len=max_draft_len), payload -def _build_exact_sps_cost_table( - payload: dict[str, object], +def _check_setup( + recorded: Mapping[str, object], live: CostTableSetup | None, *, path: str | Path +) -> None: + if live is None: + return + live_values = live.to_json() + mismatches = [ + f"{name}: table={recorded[name]!r}, runtime={live_values[name]!r}" + for name in _SETUP_ERROR_FIELDS + if name in recorded and recorded[name] != live_values[name] + ] + if mismatches: + raise ValueError( + f"DSpark cost table {path} was measured on a different setup (" + + "; ".join(mismatches) + + "). Step costs do not transfer across GPU models or parallel layouts; " + "generate a table for this setup with: " + + PROFILE_COMMAND_HINT.format(k="", path=path) + ) + recorded_version = recorded.get("tensorrt_llm_version") + if ( + recorded_version + and live.tensorrt_llm_version + and recorded_version != live.tensorrt_llm_version + ): + logger.warning( + f"DSpark cost table {path} was measured with TensorRT-LLM {recorded_version}; " + f"running {live.tensorrt_llm_version}. Costs may have drifted; regenerate the " + "table if confidence scheduling underperforms." + ) + + +def load_runtime_sps_cost_table( + path: str | Path, *, + graph_batch_sizes: Sequence[int], max_draft_len: int, + minimum_predicted_gain: float, + live_setup: CostTableSetup | None = None, ) -> ExactSpsCostTable: - exact_tables = payload["cost_tables"] - assert isinstance(exact_tables, dict) - parsed_tables: dict[int, ExactSpsCostRow] = {} - for graph_batch_size, exact_payload in exact_tables.items(): - canonical_graph_batch_size = _parse_graph_batch_size_key(graph_batch_size) - assert isinstance(exact_payload, dict) - parsed_tables[canonical_graph_batch_size] = ExactSpsCostRow( - token_counts=tuple( - _require_json_int(value, field="rank-local verifier budget", minimum=0) - for value in exact_payload["token_counts"] - ), - step_time_ms=tuple( - _require_positive_finite_number(value, field="SPS step time") - for value in exact_payload["step_time_ms"] - ), - ) - drain_guard_payload = payload.get("iteration_drain_guard") - iteration_drain_guard = None - if drain_guard_payload is not None: - assert isinstance(drain_guard_payload, dict) - iteration_drain_guard = ExactSpsDrainGuard( - loss_multiplier=_require_positive_finite_number( - drain_guard_payload["loss_multiplier"], - field="iteration drain loss_multiplier", - ), - mean_output_tokens_per_request_iteration=_require_positive_finite_number( - drain_guard_payload["mean_output_tokens_per_request_iteration"], - field="iteration drain mean_output_tokens_per_request_iteration", - ), - minimum_group_value_ms=_require_nonnegative_finite_number( - drain_guard_payload["minimum_group_value_ms"], - field="iteration drain minimum_group_value_ms", - ), - tail_graph_batch_size=_require_json_int( - drain_guard_payload["tail_graph_batch_size"], - field="iteration drain tail_graph_batch_size", - minimum=1, - ), - source_result_sha256=_require_sha256( - drain_guard_payload["source_result_sha256"], - field="iteration drain source_result_sha256", - ), + """Load the costs for ``max_draft_len`` and the captured graph ladder. + + Raises with a regenerate hint when the file is missing, has no entry for + this K, lacks a row for a captured graph batch size, or was measured on + another GPU model or parallel layout. Rows for graph sizes the engine + does not capture are ignored. + """ + payload = _read_cost_table_file(path) + rows, recorded_setup = read_cost_table_rows(payload, max_draft_len=max_draft_len, path=path) + _check_setup(recorded_setup, live_setup, path=path) + captured = sorted( + { + _require_exact_int(value, field="captured graph batch size", minimum=1) + for value in graph_batch_sizes + } + ) + if not captured: + raise ValueError("DSpark confidence scheduling requires captured CUDA-graph batch sizes") + missing = [graph_batch_size for graph_batch_size in captured if graph_batch_size not in rows] + if missing: + raise ValueError( + f"DSpark cost table {path} (K={max_draft_len}) has no measurements for captured " + f"graph batch sizes {missing}; it covers {sorted(rows)}. Either set " + "cuda_graph_config.batch_sizes to covered sizes or regenerate the table with: " + + PROFILE_COMMAND_HINT.format(k=max_draft_len, path=path) ) return ExactSpsCostTable( - tables=parsed_tables, + tables={graph_batch_size: rows[graph_batch_size] for graph_batch_size in captured}, max_draft_len=max_draft_len, - minimum_predicted_gain=_require_nonnegative_finite_number( - payload["minimum_predicted_gain"], field="minimum_predicted_gain" - ), - iteration_drain_guard=iteration_drain_guard, + minimum_predicted_gain=minimum_predicted_gain, ) -def _canonical_json_sha256(value: object) -> str: - encoded = json.dumps(value, sort_keys=True, separators=(",", ":")).encode() - return hashlib.sha256(encoded).hexdigest() +def build_cost_table_payload( + *, + setup: CostTableSetup, + costs: Mapping[int, Mapping[int, Sequence[tuple[int, float]]]], + extra: Mapping[str, object] | None = None, +) -> dict[str, object]: + """Serialize ``{K: {G: [(V, ms), ...]}}`` as a schema-3 cost table.""" + serialized: dict[str, object] = {} + for k, rows in sorted(costs.items()): + serialized_rows = {} + for graph_batch_size, cells in sorted(rows.items()): + ordered = sorted((int(v), float(ms)) for v, ms in cells) + row = ExactSpsCostRow( + token_counts=tuple(v for v, _ in ordered), + step_time_ms=tuple(ms for _, ms in ordered), + ) + serialized_rows[str(int(graph_batch_size))] = { + "token_counts": list(row.token_counts), + "step_time_ms": [round(ms, 4) for ms in row.step_time_ms], + } + # Validate the per-K grid with the same rules the runtime applies. + ExactSpsCostTable( + tables=_parse_rows(serialized_rows, where=f"K={k}"), + max_draft_len=int(k), + ) + serialized[str(int(k))] = serialized_rows + payload: dict[str, object] = { + "schema_version": COST_TABLE_SCHEMA_VERSION, + "setup": setup.to_json(), + "costs": serialized, + } + if extra: + payload.update(extra) + return payload -_V2_TOP_LEVEL_FIELDS = { - "cost_tables", - "engine_fingerprint", - "engine_fingerprint_sha256", - "measurements", - "minimum_predicted_gain", - "schema_version", -} -_V2_OPTIONAL_TOP_LEVEL_FIELDS = {"iteration_drain_guard"} -_V2_TABLE_FIELDS = {"step_time_ms", "token_counts"} -_V2_DRAIN_GUARD_FIELDS = { - "loss_multiplier", - "mean_output_tokens_per_request_iteration", - "minimum_group_value_ms", - "source_result_sha256", - "tail_graph_batch_size", -} -_V2_MEASUREMENT_FIELDS = { - "rank_local_graph_batch_size", - "rank_local_verifier_budget", - "source_result_sha256", - "step_time_ms", -} -_V2_FINGERPRINT_FIELDS = { - "global_graph_batch_sizes", - "gpu", - "gpu_count", - "gpu_snapshot_sha256", - "max_draft_len", - "rank_local_graph_batch_sizes", - "runtime_snapshot", - "source_diff_sha256", - "source_head", - "topology", -} - - -def _validate_exact_fields( - value: object, - *, - name: str, - fields: set[str], - optional_fields: set[str] | None = None, +def merge_cost_table_payloads( + existing: Mapping[str, object], new: Mapping[str, object] ) -> dict[str, object]: - if not isinstance(value, dict): - raise TypeError(f"{name} must be a JSON object") - non_string_keys = [key for key in value if not isinstance(key, str)] - if non_string_keys: - raise TypeError(f"{name} field names must be strings") - missing = sorted(fields - value.keys()) - if missing: - raise ValueError(f"{name} is missing required fields: " + ", ".join(missing)) - unknown = sorted(value.keys() - fields - (optional_fields or set())) - if unknown: - raise ValueError(f"{name} has unknown fields: " + ", ".join(unknown)) - null_fields = sorted(key for key, item in value.items() if item is None) - if null_fields: - raise ValueError(f"{name} has null fields: " + ", ".join(null_fields)) - return value + """Add or replace ``new``'s (K, G) rows in ``existing`` (same setup required). + Rows of ``existing`` that ``new`` does not measure are kept, so a table can + be extended one K or one graph batch size at a time. + """ + if existing.get("schema_version") != COST_TABLE_SCHEMA_VERSION: + raise ValueError( + "Only schema-3 cost tables can be merged; write the new costs to a new file" + ) + old_setup = dict(existing.get("setup") or {}) + new_setup = dict(new.get("setup") or {}) + conflicts = [ + name + for name in _SETUP_ERROR_FIELDS + if name in old_setup and name in new_setup and old_setup[name] != new_setup[name] + ] + if conflicts: + raise ValueError( + "Cannot merge cost tables measured on different setups: " + ", ".join(conflicts) + ) + merged = dict(existing) + merged["setup"] = new_setup or old_setup + costs = _merge_nested(existing.get("costs") or {}, new.get("costs") or {}) + merged["costs"] = { + k: dict(sorted(rows.items(), key=lambda item: int(item[0]))) + for k, rows in sorted(costs.items(), key=lambda item: int(item[0])) + } + for key, value in new.items(): + if key not in ("schema_version", "setup", "costs"): + # Per-K metadata such as the profiler's per-(G, V) sample counts. + merged[key] = _merge_nested(merged.get(key), value, depth=3) + return merged -def _validate_v2_shape(payload: dict[str, object]) -> None: - _validate_exact_fields( - payload, - name="schema-v2 SPS artifact", - fields=_V2_TOP_LEVEL_FIELDS, - optional_fields=_V2_OPTIONAL_TOP_LEVEL_FIELDS, - ) - exact_tables = payload["cost_tables"] - if not isinstance(exact_tables, dict) or not exact_tables: - raise TypeError("SPS cost_tables must be a non-empty object keyed by graph batch size") - for graph_batch_size, table_payload in exact_tables.items(): - _parse_graph_batch_size_key(graph_batch_size) - table_payload = _validate_exact_fields( - table_payload, name=f"SPS cost table for G={graph_batch_size}", fields=_V2_TABLE_FIELDS - ) - token_counts = table_payload["token_counts"] - step_time_ms = table_payload["step_time_ms"] - if not isinstance(token_counts, list) or not token_counts: - raise TypeError( - f"SPS cost table for G={graph_batch_size} token_counts must be a non-empty list" - ) - if not isinstance(step_time_ms, list) or not step_time_ms: - raise TypeError( - f"SPS cost table for G={graph_batch_size} step_time_ms must be a non-empty list" - ) - if len(token_counts) != len(step_time_ms): - raise ValueError(f"SPS cost table for G={graph_batch_size} has mismatched cells") - for value in token_counts: - _require_json_int(value, field="rank-local verifier budget", minimum=0) - for value in step_time_ms: - _require_positive_finite_number(value, field="SPS step time") - fingerprint = _validate_exact_fields( - payload["engine_fingerprint"], - name="SPS engine_fingerprint", - fields=_V2_FINGERPRINT_FIELDS, - ) - for key in ("gpu", "runtime_snapshot", "source_head", "topology"): - if not isinstance(fingerprint[key], str) or not fingerprint[key]: - raise TypeError(f"SPS engine_fingerprint {key} must be a non-empty string") - for key in ("gpu_snapshot_sha256", "source_diff_sha256"): - _require_sha256(fingerprint[key], field=f"engine_fingerprint {key}") - _require_json_int(fingerprint["gpu_count"], field="engine_fingerprint gpu_count", minimum=1) - _require_json_int( - fingerprint["max_draft_len"], field="engine_fingerprint max_draft_len", minimum=1 - ) - for key in ("rank_local_graph_batch_sizes", "global_graph_batch_sizes"): - values = fingerprint[key] - if not isinstance(values, list) or not values: - raise TypeError(f"engine_fingerprint {key} must be a non-empty list") - for value in values: - _require_json_int(value, field=f"engine_fingerprint {key}", minimum=1) - if len(set(values)) != len(values): - raise ValueError(f"engine_fingerprint {key} must not contain duplicates") - _require_sha256(payload["engine_fingerprint_sha256"], field="engine_fingerprint_sha256") - _require_nonnegative_finite_number( - payload["minimum_predicted_gain"], field="minimum_predicted_gain" - ) - drain_guard = payload.get("iteration_drain_guard") - if drain_guard is not None: - drain_guard = _validate_exact_fields( - drain_guard, - name="SPS iteration_drain_guard", - fields=_V2_DRAIN_GUARD_FIELDS, - ) - _require_positive_finite_number( - drain_guard["loss_multiplier"], - field="iteration drain loss_multiplier", - ) - _require_positive_finite_number( - drain_guard["mean_output_tokens_per_request_iteration"], - field="iteration drain mean_output_tokens_per_request_iteration", - ) - _require_nonnegative_finite_number( - drain_guard["minimum_group_value_ms"], - field="iteration drain minimum_group_value_ms", - ) - _require_json_int( - drain_guard["tail_graph_batch_size"], - field="iteration drain tail_graph_batch_size", - minimum=1, - ) - _require_sha256( - drain_guard["source_result_sha256"], - field="iteration drain source_result_sha256", - ) - measurements = payload["measurements"] - if not isinstance(measurements, list) or not measurements: - raise TypeError("SPS measurements must be a non-empty list") - for index, measurement in enumerate(measurements): - measurement = _validate_exact_fields( - measurement, - name=f"SPS measurement {index}", - fields=_V2_MEASUREMENT_FIELDS, - ) - _require_json_int( - measurement["rank_local_graph_batch_size"], - field="measurement graph batch size", - minimum=1, - ) - _require_json_int( - measurement["rank_local_verifier_budget"], - field="measurement verifier budget", - minimum=0, - ) - _require_positive_finite_number(measurement["step_time_ms"], field="measurement step time") - _require_sha256( - measurement["source_result_sha256"], field="measurement source_result_sha256" - ) + +def _merge_nested(old: object, new: object, depth: int = 2) -> object: + """``new`` over ``old``, merging dicts ``depth`` levels deep.""" + if depth == 0 or not isinstance(old, Mapping) or not isinstance(new, Mapping): + return new + merged = dict(old) + for key, value in new.items(): + merged[key] = _merge_nested(old.get(key), value, depth - 1) + return merged + + +def _canonical_json_sha256(value: object) -> str: + encoded = json.dumps(value, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(encoded).hexdigest() def _parse_graph_batch_size_key(value: object) -> int: @@ -639,12 +579,9 @@ def _parse_graph_batch_size_key(value: object) -> int: or value.startswith("0") ): raise TypeError( - f"SPS graph batch size key {value!r} must be a canonical positive integer string" + f"Cost table graph batch size key {value!r} must be a canonical positive integer string" ) - parsed = int(value) - if parsed < 1 or str(parsed) != value: - raise ValueError(f"SPS graph batch size key {value!r} must be positive and canonical") - return parsed + return int(value) def _require_json_int(value: object, *, field: str, minimum: int) -> int: @@ -682,465 +619,20 @@ def _require_nonnegative_finite_number(value: object, *, field: str) -> float: return parsed -def _require_sha256(value: object, *, field: str) -> str: - if ( - not isinstance(value, str) - or len(value) != 64 - or any(char not in "0123456789abcdef" for char in value) - ): - raise ValueError(f"{field} must be a lowercase SHA256 digest") - return value - - -def validate_sps_cost_table_payload( - payload: dict[str, object], - *, - graph_batch_sizes: Sequence[int], - max_draft_len: int, - live_engine_fingerprint: dict[str, object] | None = None, -) -> dict[str, object]: - """Validate exact schema-v2 coverage and engine provenance. - - Schema-v2 candidate budgets come directly from each per-G table. ``V=0`` - is required for every G and denotes the measured native static physical-K path; - positive V cells are arbitrary ragged candidates and need not be uniform - length multiples. Other schemas are rejected. - - Args: - payload: Parsed schema-v2 SPS artifact. - 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. - - Returns: - The artifact's authenticated engine fingerprint. - - Raises: - TypeError: If the artifact has malformed container types. - ValueError: If its schema, fingerprint, coverage, or provenance does - not exactly describe the configured graph ladder. - """ - if not isinstance(payload, dict): - raise TypeError("SPS cost artifact must contain a JSON object") - 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") - _validate_v2_shape(payload) - exact_tables = payload["cost_tables"] - fingerprint = payload["engine_fingerprint"] - assert isinstance(exact_tables, dict) - assert isinstance(fingerprint, dict) - - if live_engine_fingerprint is None: - raise ValueError( - "Schema-v2 SPS validation requires an independently supplied live engine fingerprint" - ) - if 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 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) - ) - - configured_graph_batch_sizes = { - _require_json_int(value, field="runtime graph batch size", minimum=1) - for value in graph_batch_sizes - } - if not configured_graph_batch_sizes: - raise ValueError("graph_batch_sizes must not be empty") - if len(configured_graph_batch_sizes) != len(graph_batch_sizes): - raise ValueError("runtime graph_batch_sizes must not contain duplicates") - measured_tables = { - _parse_graph_batch_size_key(graph_batch_size): table_payload - for graph_batch_size, table_payload in exact_tables.items() - } - drain_guard = payload.get("iteration_drain_guard") - if drain_guard is not None: - assert isinstance(drain_guard, dict) - tail_graph_batch_size = _require_json_int( - drain_guard["tail_graph_batch_size"], - field="iteration drain tail_graph_batch_size", - minimum=1, - ) - if tail_graph_batch_size not in measured_tables: - raise ValueError( - "Iteration drain guard requires a measured native tail table for " - f"G={tail_graph_batch_size}" - ) - fingerprint_graph_batch_sizes = set(fingerprint["rank_local_graph_batch_sizes"]) - if ( - configured_graph_batch_sizes != set(measured_tables) - or configured_graph_batch_sizes != fingerprint_graph_batch_sizes - ): - raise ValueError( - "Multi-G SPS graph batch sizes must match exactly across runtime " - "graphs, cost_tables, and engine_fingerprint" - ) - runtime_max_draft_len = _require_json_int( - max_draft_len, field="runtime max_draft_len", minimum=1 - ) - if fingerprint["max_draft_len"] != runtime_max_draft_len: - raise ValueError( - f"SPS engine_fingerprint max_draft_len does not match runtime K={max_draft_len}" - ) - - table_cells: dict[tuple[int, int], float] = {} - for graph_batch_size, table_payload in measured_tables.items(): - assert isinstance(table_payload, dict) - token_counts = [ - _require_json_int(value, field="rank-local verifier budget", minimum=0) - for value in table_payload["token_counts"] - ] - step_times = [ - _require_positive_finite_number(value, field="SPS step time") - for value in table_payload["step_time_ms"] - ] - if len(token_counts) != len(step_times): - raise ValueError(f"SPS cost table for G={graph_batch_size} has mismatched cells") - if not token_counts or any( - right <= left for left, right in zip(token_counts, token_counts[1:]) - ): - raise ValueError( - f"SPS cost table for G={graph_batch_size} budgets must be strictly increasing" - ) - if token_counts[0] != 0: - raise ValueError( - "Every schema-v2 G requires V=0 native static physical-K fallback; " - f"missing for G={graph_batch_size}" - ) - max_verifier_budget = graph_batch_size * (runtime_max_draft_len + 1) - invalid_positive_budgets = [ - verifier_tokens - for verifier_tokens in token_counts[1:] - if (verifier_tokens < graph_batch_size or verifier_tokens > max_verifier_budget) - ] - if invalid_positive_budgets: - raise ValueError( - "Schema-v2 positive verifier budgets must satisfy " - "G <= V <= G*(K+1); " - f"G={graph_batch_size}, K={runtime_max_draft_len}, " - f"invalid V={invalid_positive_budgets}" - ) - for verifier_tokens, step_time in zip(token_counts, step_times): - table_cells[(graph_batch_size, verifier_tokens)] = step_time - - measurements = payload["measurements"] - assert isinstance(measurements, list) - measurement_cells: dict[tuple[int, int], float] = {} - for item in measurements: - assert isinstance(item, dict) - pair = ( - _require_json_int( - item["rank_local_graph_batch_size"], field="measurement graph batch size", minimum=1 - ), - _require_json_int( - item["rank_local_verifier_budget"], field="measurement verifier budget", minimum=0 - ), - ) - if pair in measurement_cells: - raise ValueError(f"duplicate SPS measurement provenance for {pair}") - measurement_cells[pair] = _require_positive_finite_number( - item["step_time_ms"], field="measurement step time" - ) - if set(measurement_cells) != set(table_cells): - missing = sorted(set(table_cells) - set(measurement_cells)) - extras = sorted(set(measurement_cells) - set(table_cells)) - raise ValueError( - "SPS cost table and measurement provenance cells must match exactly; " - f"missing={missing}, extras={extras}" - ) - mismatched_cells = sorted( - pair for pair, step_time in table_cells.items() if measurement_cells[pair] != step_time - ) - if mismatched_cells: - raise ValueError( - f"SPS cost table cells do not match measurement provenance: {mismatched_cells}" - ) - - gpu_count = fingerprint["gpu_count"] - global_graph_batch_sizes = set(fingerprint["global_graph_batch_sizes"]) - if global_graph_batch_sizes != { - graph_batch_size * gpu_count for graph_batch_size in configured_graph_batch_sizes - }: - raise ValueError( - "SPS engine_fingerprint global_graph_batch_sizes are inconsistent " - "with rank-local G values and gpu_count" - ) - return fingerprint - - -def _validate_live_fingerprint_values(fingerprint: dict[str, object]) -> None: - for key in ("gpu", "runtime_snapshot", "source_head", "topology"): - if not isinstance(fingerprint[key], str) or not fingerprint[key]: - raise TypeError(f"Live engine fingerprint {key} must be a non-empty string") - for key in ("gpu_snapshot_sha256", "source_diff_sha256"): - _require_sha256(fingerprint[key], field=f"live engine fingerprint {key}") - _require_json_int( - fingerprint["gpu_count"], field="live engine fingerprint gpu_count", minimum=1 - ) - _require_json_int( - fingerprint["max_draft_len"], field="live engine fingerprint max_draft_len", minimum=1 - ) - for key in ("rank_local_graph_batch_sizes", "global_graph_batch_sizes"): - values = fingerprint[key] - if not isinstance(values, list) or not values: - raise TypeError(f"live engine fingerprint {key} must be a non-empty list") - for value in values: - _require_json_int(value, field=f"live engine fingerprint {key}", minimum=1) - if len(set(values)) != len(values): - raise ValueError(f"live engine fingerprint {key} must not contain duplicates") - - -def evaluate_confidence_workload_admission( - *, - policy_steps: int, - compact_choices: int, - gross_compact_value_ms_lower_bound: float, - fixed_confidence_path_overhead_ms_upper_bound: float, - safety_margin_ms: float, -) -> tuple[bool, float]: - """Evaluate aggregate confidence economics without an endogenous ratio. - - ``compact_choices`` comes from replaying the exact marginal selector - without charging fixed confidence-path overhead to one candidate. The - candidate set therefore cannot depend on ``policy_steps / compact_choices``. - """ - policy_steps = _require_exact_int(policy_steps, field="policy_steps", minimum=1) - compact_choices = _require_exact_int(compact_choices, field="compact_choices", minimum=0) - if compact_choices > policy_steps: - raise ValueError( - f"compact_choices must not exceed policy_steps; got {compact_choices} > {policy_steps}" - ) - gross_value_ms = _require_nonnegative_finite_number( - gross_compact_value_ms_lower_bound, field="gross compact value lower bound" - ) - overhead_ms = _require_nonnegative_finite_number( - fixed_confidence_path_overhead_ms_upper_bound, - field="fixed confidence-path overhead upper bound", - ) - margin_ms = _require_nonnegative_finite_number( - safety_margin_ms, field="confidence admission safety margin" - ) - net_value_ms = gross_value_ms - policy_steps * overhead_ms - margin_ms - return compact_choices > 0 and net_value_ms > 0.0, net_value_ms - - -_WORKLOAD_ADMISSION_FIELDS = { - "admitted", - "calibration_result_sha256", - "compact_choices", - "fixed_confidence_path_overhead_ms_upper_bound", - "gross_compact_value_ms_lower_bound", - "live_engine_fingerprint_sha256", - "physical_k", - "policy_steps", - "runtime_snapshot", - "safety_margin_ms", - "selector_identity_sha256", - "selector_replay_sha256", - "source_diff_sha256", - "source_head", - "sps_cost_table_sha256", - "workload_identity_sha256", -} -_WORKLOAD_ADMISSION_ENVELOPE_FIELDS = { - "admission", - "admission_sha256", - "schema_version", -} - - -def _read_pinned_bytes(path: str | Path) -> tuple[bytes, str]: - data = Path(path).read_bytes() - return data, hashlib.sha256(data).hexdigest() - - -def _parse_strict_json_object(data: bytes, *, name: str) -> dict[str, object]: - def reject_duplicate_fields(pairs: list[tuple[str, object]]) -> dict[str, object]: - parsed = {} - for key, value in pairs: - if key in parsed: - raise ValueError(f"{name} has duplicate field {key!r}") - parsed[key] = value - return parsed - - payload = json.loads(data.decode("utf-8"), object_pairs_hook=reject_duplicate_fields) - if not isinstance(payload, dict): - raise TypeError(f"{name} must contain a JSON object") - return payload - - -def load_confidence_workload_admission( - path: str | Path, - *, - expected_receipt_sha256: str, - sps_cost_table_path: str | Path, - live_engine_fingerprint_path: str | Path, - physical_k: int, -) -> ConfidenceWorkloadAdmission: - """Authenticate and evaluate one deployment-level confidence receipt. - - The caller supplies a trusted receipt digest and immutable artifact files. - The receipt binds the cost table, live fingerprint and selector identity; - its workload, calibration and replay digests identify externally qualified - evidence. This loader does not generate or independently qualify that evidence. - """ - physical_k = _require_exact_int(physical_k, field="physical_k", minimum=1) - expected_receipt_sha256 = _require_sha256( - expected_receipt_sha256, - field="expected confidence workload admission receipt SHA256", - ) - receipt_bytes, receipt_sha256 = _read_pinned_bytes(path) - if receipt_sha256 != expected_receipt_sha256: - raise ValueError("Confidence workload admission receipt does not match its pinned SHA256") - envelope = _validate_exact_fields( - _parse_strict_json_object(receipt_bytes, name="confidence workload admission receipt"), - name="confidence workload admission receipt", - fields=_WORKLOAD_ADMISSION_ENVELOPE_FIELDS, - ) - if type(envelope["schema_version"]) is not int or envelope["schema_version"] != 1: - raise ValueError("Confidence workload admission receipts require schema_version=1") - admission = _validate_exact_fields( - envelope["admission"], - name="confidence workload admission body", - fields=_WORKLOAD_ADMISSION_FIELDS, - ) - admission_sha256 = _require_sha256( - envelope["admission_sha256"], field="confidence workload admission SHA256" - ) - if admission_sha256 != _canonical_json_sha256(admission): - raise ValueError("Confidence workload admission SHA256 does not match its body") - if type(admission["admitted"]) is not bool: - raise TypeError("confidence workload admitted must be a JSON boolean") - receipt_physical_k = _require_json_int( - admission["physical_k"], field="confidence workload physical_k", minimum=1 - ) - if receipt_physical_k != physical_k: - raise ValueError( - "Confidence workload admission physical_k does not match the configured " - f"physical K: {receipt_physical_k} != {physical_k}" - ) - - digest_fields = ( - "calibration_result_sha256", - "live_engine_fingerprint_sha256", - "selector_identity_sha256", - "selector_replay_sha256", - "source_diff_sha256", - "sps_cost_table_sha256", - "workload_identity_sha256", - ) - for field_name in digest_fields: - _require_sha256(admission[field_name], field=f"confidence workload {field_name}") - if admission["selector_identity_sha256"] != EXACT_SPS_SELECTOR_IDENTITY_SHA256: - raise ValueError("Confidence workload admission selector identity is stale or unknown") - sps_bytes, sps_sha256 = _read_pinned_bytes(sps_cost_table_path) - if admission["sps_cost_table_sha256"] != sps_sha256: - raise ValueError("Confidence workload admission SPS cost-table SHA256 does not match") - live_fingerprint_bytes, live_fingerprint_sha256 = _read_pinned_bytes( - live_engine_fingerprint_path - ) - if admission["live_engine_fingerprint_sha256"] != live_fingerprint_sha256: - raise ValueError( - "Confidence workload admission live-engine fingerprint SHA256 does not match" - ) - - live_fingerprint = _validate_exact_fields( - _parse_strict_json_object(live_fingerprint_bytes, name="live engine fingerprint"), - name="live engine fingerprint", - fields=_V2_FINGERPRINT_FIELDS, - ) - _validate_live_fingerprint_values(live_fingerprint) - if live_fingerprint["max_draft_len"] != physical_k: - raise ValueError( - "Live engine fingerprint max_draft_len does not match confidence workload " - f"physical K: {live_fingerprint['max_draft_len']} != {physical_k}" - ) - for field_name in ("runtime_snapshot", "source_diff_sha256", "source_head"): - if admission[field_name] != live_fingerprint[field_name]: - raise ValueError( - f"Confidence workload admission {field_name} does not match the live engine" - ) - for field_name in ("runtime_snapshot", "source_head"): - if not isinstance(admission[field_name], str) or not admission[field_name]: - raise TypeError(f"confidence workload {field_name} must be a non-empty string") - - sps_payload = _parse_strict_json_object(sps_bytes, name="SPS cost artifact") - validate_sps_cost_table_payload( - sps_payload, - graph_batch_sizes=live_fingerprint["rank_local_graph_batch_sizes"], - max_draft_len=physical_k, - live_engine_fingerprint=live_fingerprint, - ) - _build_exact_sps_cost_table(sps_payload, max_draft_len=physical_k) - admitted, net_value_ms = evaluate_confidence_workload_admission( - policy_steps=admission["policy_steps"], - compact_choices=admission["compact_choices"], - gross_compact_value_ms_lower_bound=admission["gross_compact_value_ms_lower_bound"], - fixed_confidence_path_overhead_ms_upper_bound=admission[ - "fixed_confidence_path_overhead_ms_upper_bound" - ], - safety_margin_ms=admission["safety_margin_ms"], - ) - if admission["admitted"] is not admitted: - raise ValueError( - "Confidence workload admission decision does not match the aggregate economics" - ) - return ConfidenceWorkloadAdmission( - admitted=admitted, - physical_k=physical_k, - policy_steps=int(admission["policy_steps"]), - compact_choices=int(admission["compact_choices"]), - gross_compact_value_ms_lower_bound=float(admission["gross_compact_value_ms_lower_bound"]), - fixed_confidence_path_overhead_ms_upper_bound=float( - admission["fixed_confidence_path_overhead_ms_upper_bound"] - ), - safety_margin_ms=float(admission["safety_margin_ms"]), - net_value_ms_lower_bound=net_value_ms, - admission_identity_sha256=admission_sha256, - receipt_sha256=receipt_sha256, - ) - - def select_exact_sps_candidate( *, graph_batch_size: int, native_expected_yield: float, compact_expected_yields: dict[int, float], - compact_max_yield_losses_per_request: dict[int, float], cost_table: ExactSpsCostTable, ) -> int: - """Choose a measured V, retaining native V=0 on weak group-E2E value. - - 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 - 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 - allgather as the aggregate yields. A selected tier must satisfy, strictly:: - - T(G, 0) - T(G, V) - - loss_multiplier * max_rank((Y_native - Y_V) / N_real) - * T(tail_G, 0) / mean_output_tokens_per_request_iteration - > minimum_group_value_ms + """Choose a measured V by expected goodput, keeping native V=0 by default. + + The caller supplies globally agreed expected yields after the allgather; + this pure layer only compares measured cells. A compact tier wins only + when its expected tokens per millisecond beat native by at least + ``cost_table.minimum_predicted_gain``; among winners the highest goodput + is chosen. """ graph_batch_size = _require_exact_int(graph_batch_size, field="graph_batch_size", minimum=1) measured_budgets = cost_table.production_candidate_budgets(graph_batch_size) @@ -1155,28 +647,12 @@ def select_exact_sps_candidate( raise ValueError( f"compact_expected_yields must cover every positive measured V for G={graph_batch_size}" ) - if any(type(value) is not int for value in compact_max_yield_losses_per_request): - raise TypeError("compact_max_yield_losses_per_request keys must be integer V values") - if 0 in compact_max_yield_losses_per_request: - raise ValueError("compact_max_yield_losses_per_request must not include native V=0") - if set(compact_max_yield_losses_per_request) != set(measured_budgets): - raise ValueError( - "compact_max_yield_losses_per_request must cover every positive " - f"measured V for G={graph_batch_size}" - ) compact_yields = { budget: _require_nonnegative_finite_number( compact_expected_yields[budget], field=f"expected yield for compact V={budget}" ) for budget in measured_budgets } - max_yield_losses_per_request = { - budget: _require_nonnegative_finite_number( - compact_max_yield_losses_per_request[budget], - field=(f"maximum expected yield loss per request for compact V={budget}"), - ) - for budget in measured_budgets - } native_yield = _require_nonnegative_finite_number( native_expected_yield, field="native_expected_yield" ) @@ -1187,8 +663,6 @@ def select_exact_sps_candidate( budget: compact_yields[budget] / cost_table.step_time(budget, graph_batch_size) for budget in measured_budgets } - if not compact_scores: - return 0 required_score = native_score * (1.0 + cost_table.minimum_predicted_gain) goodput_budgets = [ budget @@ -1197,36 +671,4 @@ def select_exact_sps_candidate( ] if not goodput_budgets: 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 - 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 = {} - for budget in goodput_budgets: - compact_step_ms = cost_table.step_time(budget, graph_batch_size) - group_value_ms = ( - native_step_ms - - compact_step_ms - - drain_guard.loss_multiplier - * max_yield_losses_per_request[budget] - * tail_step_ms - / drain_guard.mean_output_tokens_per_request_iteration - ) - if group_value_ms > drain_guard.minimum_group_value_ms: - group_values_ms[budget] = group_value_ms - if not group_values_ms: - return 0 - return int( - max( - group_values_ms, - key=lambda budget: ( - group_values_ms[budget], - compact_scores[budget], - budget, - ), - ) - ) + return int(max(goodput_budgets, key=lambda budget: (compact_scores[budget], budget))) diff --git a/tensorrt_llm/_torch/speculative/dspark_profile.py b/tensorrt_llm/_torch/speculative/dspark_profile.py new file mode 100644 index 000000000000..1c61901fda30 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_profile.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Runtime side of ``tensorrt_llm.tools.dspark_profile_costs``. + +When ``TLLM_DSPARK_PROFILE_COSTS`` names an output file, the executor replaces +the confidence policy's choice in full decode-only steps with a fixed +round-robin over the captured verifier budgets ``V`` (native first) and records +the wall-clock period between consecutive scheduling decisions. Each budget is +held for several consecutive steps so the overlap scheduler reaches steady +state, and the first steps of every hold are discarded. Every rank runs the +same schedule because it depends only on globally agreed state; rank 0 writes +the samples. +""" + +import json +import os +import time +from collections import defaultdict +from pathlib import Path +from typing import Optional, Sequence + +PROFILE_OUTPUT_ENV = "TLLM_DSPARK_PROFILE_COSTS" +PROFILE_HOLD_ENV = "TLLM_DSPARK_PROFILE_HOLD_STEPS" +PROFILE_WARMUP_ENV = "TLLM_DSPARK_PROFILE_WARMUP_STEPS" + + +class DSparkCostProfiler: + """Force verifier budgets round-robin and record whole-step periods.""" + + def __init__( + self, + output_path: str, + *, + write_samples: bool, + hold_steps: int = 16, + warmup_steps: int = 3, + ) -> None: + if hold_steps <= warmup_steps: + raise ValueError("hold_steps must exceed warmup_steps") + self.output_path = Path(output_path) + self.write_samples = write_samples + self.hold_steps = hold_steps + self.warmup_steps = warmup_steps + self._decisions_per_g: dict[int, int] = defaultdict(int) + self._samples: dict[tuple[int, int], list[float]] = defaultdict(list) + self._previous: Optional[tuple[int, int, int, float, int]] = None + self._unflushed = 0 + + @classmethod + def from_env(cls, *, rank: int) -> Optional["DSparkCostProfiler"]: + output_path = os.environ.get(PROFILE_OUTPUT_ENV) + if not output_path: + return None + return cls( + output_path, + write_samples=rank == 0, + hold_steps=int(os.environ.get(PROFILE_HOLD_ENV, "16")), + warmup_steps=int(os.environ.get(PROFILE_WARMUP_ENV, "3")), + ) + + def choose( + self, + *, + iteration: int, + graph_batch_size: int, + budgets: Sequence[int], + all_ranks_full: bool, + ) -> int: + """Return the forced V (0 = native) for this decode step. + + ``budgets`` lists native 0 first, then the captured compact budgets of + ``graph_batch_size``. Only steps in which every rank has exactly + ``graph_batch_size`` real requests are timed. + """ + now = time.perf_counter() + previous = self._previous + if previous is not None and iteration == previous[0] + 1: + _, prev_g, prev_v, prev_time, prev_position = previous + if prev_position >= self.warmup_steps: + self._samples[(prev_g, prev_v)].append((now - prev_time) * 1000.0) + self._unflushed += 1 + if not all_ranks_full: + # The wave is filling or draining: persist what the full steps + # measured so far, since the run may end soon. + self._previous = None + self._maybe_flush(force=True) + return 0 + count = self._decisions_per_g[graph_batch_size] + self._decisions_per_g[graph_batch_size] = count + 1 + cell = (count // self.hold_steps) % len(budgets) + budget = int(budgets[cell]) + self._previous = (iteration, graph_batch_size, budget, now, count % self.hold_steps) + self._maybe_flush(force=count % self.hold_steps == 0) + return budget + + def _maybe_flush(self, *, force: bool) -> None: + if not self.write_samples or self._unflushed == 0: + return + if not force and self._unflushed < self.hold_steps: + return + payload = { + "samples": [ + {"graph_batch_size": g, "verifier_budget": v, "step_time_ms": times} + for (g, v), times in sorted(self._samples.items()) + ] + } + tmp_path = self.output_path.with_name(self.output_path.name + ".tmp") + tmp_path.write_text(json.dumps(payload)) + os.replace(tmp_path, self.output_path) + self._unflushed = 0 + + def flush(self) -> None: + self._maybe_flush(force=True) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 11dc8d41aa7d..120bbd97d4b3 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3793,9 +3793,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``. @@ -3837,6 +3836,46 @@ class DSparkDecodingConfig(DecodingBaseConfig): "read from the draft model config (dspark_markov_head_type), " "defaulting to \"vanilla\".") + enable_confidence_scheduling: bool = Field( + default=False, + description= + "Enable confidence-guided selection of a measured verifier-token budget " + "for decode-only batches. The full physical draft block is still " + "proposed; target verification and acceptance rules are unchanged. " + "Requires trained confidence-head weights and a cost table " + "(confidence_cost_table_path). Native verification is kept whenever no " + "compact budget is predicted to win by confidence_min_predicted_gain.", + status="prototype") + + confidence_cost_table_path: Optional[str] = Field( + default=None, + description= + "Path to a DSpark cost table: measured decode step times per physical " + "K, rank-local CUDA-graph batch size and verifier-token budget. One " + "file may hold several K values; the entry for max_draft_len is used. " + "Generate one with `python -m tensorrt_llm.tools.dspark_profile_costs`. " + "Required when enable_confidence_scheduling=True.", + status="prototype") + + confidence_calibration_path: Optional[str] = Field( + default=None, + description= + "Path to a JSON file with per-position temperature-scaling values " + "(key 'sts_temperatures' or 'temperatures') used to calibrate the " + "confidence head. A vector longer than max_draft_len is truncated to " + "its first max_draft_len entries, so one file serves every K. If None, " + "the raw sigmoid is used.", + status="prototype") + + confidence_min_predicted_gain: NonNegativeFloat = Field( + default=0.04, + description= + "Minimum predicted relative goodput gain (expected accepted tokens per " + "millisecond) a compact verifier budget needs over native " + "verification before it is selected. Larger values keep native " + "verification more often.", + status="prototype") + decoding_type: Literal["DSpark"] = Field(default="DSpark") attention_backend: Literal["AUTO", "VANILLA", "TRTLLM", "CUTEDSL"] = Field( @@ -3853,6 +3892,32 @@ class DSparkDecodingConfig(DecodingBaseConfig): "decode pass over context and block, on SM100/SM103/SM107. Which " "kernel each name selects: MLADSparkForCausalLM.") + @model_validator(mode="after") + def validate_confidence_scheduling(self): + if not self.enable_confidence_scheduling: + # A static config copied from a dynamic one may keep the + # confidence paths; they are simply unused. + 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. + logger.warning( + "DSpark confidence scheduling has no compact tier below the " + "native K1 layout; using the feature-off static K1 path.") + self.enable_confidence_scheduling = False + return self + if not self.confidence_cost_table_path: + raise ValueError( + "DSpark confidence scheduling requires confidence_cost_table_path. " + "Generate a cost table with `python -m " + "tensorrt_llm.tools.dspark_profile_costs` or point to an " + "existing one; graph coverage and K are checked at engine " + "start.") + return self + @model_validator(mode="after") def set_max_total_draft_tokens(self): self.max_total_draft_tokens = self.max_draft_len @@ -7401,6 +7466,122 @@ 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.") + speculative_config = self.speculative_config + if (getattr(speculative_config, "acceptance_rate_window_size", None) + and getattr(speculative_config, "acceptance_rate_threshold", + None) is not None): + raise ValueError( + "DSpark confidence scheduling does not support the acceptance-" + "rate gate (acceptance_rate_window_size / " + "acceptance_rate_threshold): once the gate disables " + "speculation, the last step's verify windows would remain in " + "use, and the gate measures acceptance over all drafted " + "tokens rather than the verified ones. Unset both fields or " + "disable confidence scheduling.") + cuda_graph_config = self.cuda_graph_config + if self.context_parallel_size > 1: + raise ValueError( + "DSpark confidence scheduling requires context_parallel_size=1; " + "the exact-shape policy does not coordinate CP groups.") + 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.") + if getattr(self.moe_config.load_balancer, "mode", + None) == "per_layer_dynamic": + raise ValueError( + "DSpark confidence scheduling does not support per-layer dynamic " + "EPLB because it disables the required CUDA graphs.") + + 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.") + @model_validator(mode="after") def normalize_prefill_cuda_graph_config(self) -> 'TorchLlmArgs': """Normalize legacy piecewise CUDA graph options into prefill fields.""" @@ -7888,6 +8069,8 @@ 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._validate_dspark_confidence_environment() if isinstance(self.speculative_config, SADecodingConfig): pool_size = self.speculative_config.global_pool_size diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 60e2051886a4..45790fca3196 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -975,7 +975,7 @@ }, "DSparkDecodingConfig": { "additionalProperties": false, - "description": "Configuration for DSpark speculative decoding.\n\nDSpark (DeepSeek) is a target-dependent, \"semi-parallel\" speculative\ndecoding method. Like DFlash it captures hidden states from several target\nlayers as cross-attention context and drafts a whole block in a single\nbackbone forward, but it additionally refines the per-position draft logits\nwith a lightweight sequential head (a low-rank Markov head, optionally an RNN\nhead). A confidence head can score each draft position for verification\nscheduling without truncating the proposed block. This configuration selects\nnative full-block drafting and does not expose confidence-scheduling or\ncalibration controls.\n\nKey features:\n- Target-dependent: captures hidden states from ``target_layer_ids``.\n- Semi-parallel: one block backbone forward + cheap sequential head refine.\n- Full-block proposals: confidence scoring does not change the draft tokens\n or the standard target-verification acceptance rule.\n\nReference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec).", + "description": "Configuration for DSpark speculative decoding.\n\nDSpark (DeepSeek) is a target-dependent, \"semi-parallel\" speculative\ndecoding method. Like DFlash it captures hidden states from several target\nlayers as cross-attention context and drafts a whole block in a single\nbackbone forward, but it additionally refines the per-position draft logits\nwith a lightweight sequential head (a low-rank Markov head, optionally an RNN\nhead). A confidence head can score each draft position for verification\nscheduling without truncating the proposed block. Confidence scheduling is\nopt-in; the default configuration retains native full-block verification.\n\nKey features:\n- Target-dependent: captures hidden states from ``target_layer_ids``.\n- Semi-parallel: one block backbone forward + cheap sequential head refine.\n- Full-block proposals: confidence scoring does not change the draft tokens\n or the standard target-verification acceptance rule.\n\nReference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec).", "properties": { "acceptance_rate_threshold": { "anyOf": [ @@ -1044,6 +1044,42 @@ "description": "Number of draft positions produced per block. If None, read from the draft model config (dspark_block_size). Should equal max_draft_len.", "title": "Block Size" }, + "confidence_calibration_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to a JSON file with per-position temperature-scaling values (key 'sts_temperatures' or 'temperatures') used to calibrate the confidence head. A vector longer than max_draft_len is truncated to its first max_draft_len entries, so one file serves every K. If None, the raw sigmoid is used.", + "status": "prototype", + "title": "Confidence Calibration Path" + }, + "confidence_cost_table_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to a DSpark cost table: measured decode step times per physical K, rank-local CUDA-graph batch size and verifier-token budget. One file may hold several K values; the entry for max_draft_len is used. Generate one with `python -m tensorrt_llm.tools.dspark_profile_costs`. Required when enable_confidence_scheduling=True.", + "status": "prototype", + "title": "Confidence Cost Table Path" + }, + "confidence_min_predicted_gain": { + "default": 0.04, + "description": "Minimum predicted relative goodput gain (expected accepted tokens per millisecond) a compact verifier budget needs over native verification before it is selected. Larger values keep native verification more often.", + "minimum": 0, + "status": "prototype", + "title": "Confidence Min Predicted Gain", + "type": "number" + }, "decoding_type": { "const": "DSpark", "default": "DSpark", @@ -1066,6 +1102,13 @@ "description": "Developer interface: dynamically adjust draft length based on active batch size in runtime.Maps batch size to draft lengths.For example: draft_len_schedule = {4:4, 8:2, 32:1} - Batch sizes 1-4: use draft_len=4 - Batch sizes 5-8: use draft_len=2 - Batch sizes 9-32: use draft_len=1 - Batch sizes 33+: use draft_len=0 (implicit, speculation disabled). Mutually exclusive with max_concurrency since draft_len_schedule implicitly support max concurrency control.", "title": "Draft Len Schedule" }, + "enable_confidence_scheduling": { + "default": false, + "description": "Enable confidence-guided selection of a measured verifier-token budget for decode-only batches. The full physical draft block is still proposed; target verification and acceptance rules are unchanged. Requires trained confidence-head weights and a cost table (confidence_cost_table_path). Native verification is kept whenever no compact budget is predicted to win by confidence_min_predicted_gain.", + "status": "prototype", + "title": "Enable Confidence Scheduling", + "type": "boolean" + }, "enable_penalty": { "default": false, "description": "If true, enables the occurrence penalties (repetition / presence / frequency) for one-model speculative decoding. Off by default because the penalties need a [num_seq_slots, vocab_size] occurrence-count workspace that is allocated up front (CUDA graphs capture fixed buffer addresses). While off, a request that asks for any of these penalties is rejected at admission rather than silently decoded without them.", diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index 5305b35e0f4a..165ab53cb100 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -975,7 +975,7 @@ }, "DSparkDecodingConfig": { "additionalProperties": false, - "description": "Configuration for DSpark speculative decoding.\n\nDSpark (DeepSeek) is a target-dependent, \"semi-parallel\" speculative\ndecoding method. Like DFlash it captures hidden states from several target\nlayers as cross-attention context and drafts a whole block in a single\nbackbone forward, but it additionally refines the per-position draft logits\nwith a lightweight sequential head (a low-rank Markov head, optionally an RNN\nhead). A confidence head can score each draft position for verification\nscheduling without truncating the proposed block. This configuration selects\nnative full-block drafting and does not expose confidence-scheduling or\ncalibration controls.\n\nKey features:\n- Target-dependent: captures hidden states from ``target_layer_ids``.\n- Semi-parallel: one block backbone forward + cheap sequential head refine.\n- Full-block proposals: confidence scoring does not change the draft tokens\n or the standard target-verification acceptance rule.\n\nReference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec).", + "description": "Configuration for DSpark speculative decoding.\n\nDSpark (DeepSeek) is a target-dependent, \"semi-parallel\" speculative\ndecoding method. Like DFlash it captures hidden states from several target\nlayers as cross-attention context and drafts a whole block in a single\nbackbone forward, but it additionally refines the per-position draft logits\nwith a lightweight sequential head (a low-rank Markov head, optionally an RNN\nhead). A confidence head can score each draft position for verification\nscheduling without truncating the proposed block. Confidence scheduling is\nopt-in; the default configuration retains native full-block verification.\n\nKey features:\n- Target-dependent: captures hidden states from ``target_layer_ids``.\n- Semi-parallel: one block backbone forward + cheap sequential head refine.\n- Full-block proposals: confidence scoring does not change the draft tokens\n or the standard target-verification acceptance rule.\n\nReference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec).", "properties": { "acceptance_rate_threshold": { "anyOf": [ @@ -1044,6 +1044,42 @@ "description": "Number of draft positions produced per block. If None, read from the draft model config (dspark_block_size). Should equal max_draft_len.", "title": "Block Size" }, + "confidence_calibration_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to a JSON file with per-position temperature-scaling values (key 'sts_temperatures' or 'temperatures') used to calibrate the confidence head. A vector longer than max_draft_len is truncated to its first max_draft_len entries, so one file serves every K. If None, the raw sigmoid is used.", + "status": "prototype", + "title": "Confidence Calibration Path" + }, + "confidence_cost_table_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to a DSpark cost table: measured decode step times per physical K, rank-local CUDA-graph batch size and verifier-token budget. One file may hold several K values; the entry for max_draft_len is used. Generate one with `python -m tensorrt_llm.tools.dspark_profile_costs`. Required when enable_confidence_scheduling=True.", + "status": "prototype", + "title": "Confidence Cost Table Path" + }, + "confidence_min_predicted_gain": { + "default": 0.04, + "description": "Minimum predicted relative goodput gain (expected accepted tokens per millisecond) a compact verifier budget needs over native verification before it is selected. Larger values keep native verification more often.", + "minimum": 0, + "status": "prototype", + "title": "Confidence Min Predicted Gain", + "type": "number" + }, "decoding_type": { "const": "DSpark", "default": "DSpark", @@ -1066,6 +1102,13 @@ "description": "Developer interface: dynamically adjust draft length based on active batch size in runtime.Maps batch size to draft lengths.For example: draft_len_schedule = {4:4, 8:2, 32:1} - Batch sizes 1-4: use draft_len=4 - Batch sizes 5-8: use draft_len=2 - Batch sizes 9-32: use draft_len=1 - Batch sizes 33+: use draft_len=0 (implicit, speculation disabled). Mutually exclusive with max_concurrency since draft_len_schedule implicitly support max concurrency control.", "title": "Draft Len Schedule" }, + "enable_confidence_scheduling": { + "default": false, + "description": "Enable confidence-guided selection of a measured verifier-token budget for decode-only batches. The full physical draft block is still proposed; target verification and acceptance rules are unchanged. Requires trained confidence-head weights and a cost table (confidence_cost_table_path). Native verification is kept whenever no compact budget is predicted to win by confidence_min_predicted_gain.", + "status": "prototype", + "title": "Enable Confidence Scheduling", + "type": "boolean" + }, "enable_penalty": { "default": false, "description": "If true, enables the occurrence penalties (repetition / presence / frequency) for one-model speculative decoding. Off by default because the penalties need a [num_seq_slots, vocab_size] occurrence-count workspace that is allocated up front (CUDA graphs capture fixed buffer addresses). While off, a request that asks for any of these penalties is rejected at admission rather than silently decoded without them.", diff --git a/tensorrt_llm/tools/dspark_profile_costs.py b/tensorrt_llm/tools/dspark_profile_costs.py new file mode 100644 index 000000000000..f687125b3703 --- /dev/null +++ b/tensorrt_llm/tools/dspark_profile_costs.py @@ -0,0 +1,440 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Generate a DSpark cost table for confidence scheduling. + +Run it on the GPUs and with the runtime YAML you will serve with:: + + python -m tensorrt_llm.tools.dspark_profile_costs \ + --model /path/to/DeepSeek-V4-Pro-DSpark --config runtime.yaml \ + --k 3 5 --output dspark-costs.json + +For each requested physical draft length K the tool loads the model once (in +its own process), then for every captured CUDA-graph batch size G runs a wave +of ``G x attention-DP size`` requests. During the wave's full decode steps the +executor cycles the verifier-token budget V through native and each compact +tier and times the whole step; another wave runs while a tier has too few +samples. The medians become the table. An existing +output file is updated in place: entries for the profiled K are replaced and +other K entries are kept. + +Point ``speculative_config.confidence_cost_table_path`` at the output and set +``enable_confidence_scheduling: true`` to use it. +""" + +import argparse +import json +import math +import os +import random +import statistics +import subprocess # nosec B404 - runs sys.executable workers with a fixed argv, no shell +import sys +import tempfile +from pathlib import Path + +import yaml + +DEFAULT_TIER_FRACTIONS = (0.5, 0.6, 0.7, 0.8) +MAX_COMPACT_CELLS_PER_G = 4 +MAX_COMPACT_CELLS_TOTAL = 32 + + +def _budget(graph_batch_size: int, k: int, fraction: float) -> int: + return int(math.floor(graph_batch_size * (1.0 + fraction * k) + 0.5)) + + +def compact_budgets(graph_batch_size: int, k: int, fractions) -> list[int]: + """Verifier budgets that verify ``fraction * K`` draft tokens per request on average. + + Budgets count the anchor row, so native is ``G * (K + 1)``. Every request + verifies at least one draft token, so a full batch needs at least ``2 * G``. + With the default fractions, K5 at G128 gives V = 448, 512, 576, 640. + """ + native = graph_batch_size * (k + 1) + budgets = {_budget(graph_batch_size, k, fraction) for fraction in fractions} + budgets = {budget for budget in budgets if 2 * graph_batch_size <= budget < native} + return sorted(budgets)[-MAX_COMPACT_CELLS_PER_G:] + + +def tiers_below_one_draft_token(ladder, k: int, fractions) -> list[tuple[int, float, int]]: + """``(G, fraction, V)`` whose budget leaves a request in a full batch with no draft token.""" + return [ + (graph_batch_size, fraction, _budget(graph_batch_size, k, fraction)) + for graph_batch_size in ladder + for fraction in fractions + if _budget(graph_batch_size, k, fraction) < 2 * graph_batch_size + ] + + +def plan_grid(ladder, k: int, fractions) -> dict[int, list[int]]: + """Compact budgets per G, giving the largest G priority within the runtime cell limit.""" + grid: dict[int, list[int]] = {} + remaining = MAX_COMPACT_CELLS_TOTAL + for graph_batch_size in sorted(ladder, reverse=True): + budgets = compact_budgets(graph_batch_size, k, fractions) if k > 1 else [] + budgets = budgets[-remaining:] if remaining > 0 else [] + remaining -= len(budgets) + grid[graph_batch_size] = budgets + return dict(sorted(grid.items())) + + +def resolve_ladder(config: dict, override) -> list[int]: + if override: + return sorted(set(int(value) for value in override)) + from tensorrt_llm.llmapi.llm_args import CudaGraphConfig + + cuda_graph = CudaGraphConfig(**(config.get("cuda_graph_config") or {})) + return sorted(set(cuda_graph.batch_sizes)) + + +def with_graph_ladder(config: dict, ladder) -> dict: + """``config`` capturing exactly the decode graphs in ``ladder``. + + The runtime needs a cost-table row for every captured graph, and the + placeholder table only has rows for the graphs being profiled. + """ + cuda_graph = dict(config.get("cuda_graph_config") or {}) + cuda_graph.update(batch_sizes=list(ladder), max_batch_size=max(ladder)) + return dict(config, cuda_graph_config=cuda_graph) + + +def without_adp_balance(config: dict) -> dict: + """``config`` with attention-DP context balancing turned off. + + Balancing holds context requests back for several iterations. With long + prompts a wave then prefills for so long that its first requests finish + before the last ones start decoding, and no step has every rank full. + It does not change the cost of a decode step. + """ + adp_config = config.get("attention_dp_config") + if not adp_config or not adp_config.get("enable_balance"): + return config + return dict(config, attention_dp_config=dict(adp_config, enable_balance=False)) + + +def short_cells( + samples_path: Path, graph_batch_size: int, budgets, min_samples: int +) -> dict[int, int]: + """Sample counts of the budgets of ``graph_batch_size`` still below ``min_samples``.""" + counts = {int(budget): 0 for budget in budgets} + if samples_path.exists(): + for entry in json.loads(samples_path.read_text())["samples"]: + budget = int(entry["verifier_budget"]) + if int(entry["graph_batch_size"]) == graph_batch_size and budget in counts: + counts[budget] = len(entry["step_time_ms"]) + return {budget: count for budget, count in counts.items() if count < min_samples} + + +def load_prompts(dataset: str | None, prompt_len: int, vocab_hint: int = 32000): + """Token-id prompts: from a trtllm-bench JSONL dataset if given, else random ids.""" + if dataset: + prompts = [] + with open(dataset, encoding="utf-8") as file: + for line in file: + record = json.loads(line) + ids = record.get("input_ids") or record.get("prompt_token_ids") + if ids: + prompts.append([int(token) for token in ids]) + if not prompts: + raise ValueError(f"{dataset} has no 'input_ids' lists") + return prompts + rng = random.Random(0) + return [[rng.randrange(100, vocab_hint) for _ in range(prompt_len)] for _ in range(256)] + + +def _profile_k(args) -> None: + """Worker: profile one K and write a single-K cost table to ``args.worker_output``.""" + import torch + + from tensorrt_llm import LLM, SamplingParams + from tensorrt_llm._torch.speculative.dspark_planner import ( + CostTableSetup, + build_cost_table_payload, + ) + from tensorrt_llm._torch.speculative.dspark_profile import PROFILE_OUTPUT_ENV + from tensorrt_llm.llmapi.llm_args import update_llm_args_with_extra_dict + from tensorrt_llm.version import __version__ as trtllm_version + + k = args.k[0] + with open(args.config, encoding="utf-8") as file: + config = yaml.safe_load(file) or {} + spec = dict(config.get("speculative_config") or {}) + if spec.get("decoding_type") != "DSpark": + raise ValueError("--config must use speculative_config.decoding_type: DSpark") + ladder = resolve_ladder(config, args.graph_batch_sizes) + grid = plan_grid(ladder, k, args.tier_fractions) + + workdir = Path(args.workdir) + grid_path = workdir / f"grid-k{k}.json" + # Placeholder prices: they only declare which (G, V) graphs to capture; + # the profiler replaces the policy while profiling. + grid_path.write_text( + json.dumps( + { + "schema_version": 3, + "costs": { + str(k): { + str(g): { + "token_counts": [0] + budgets, + "step_time_ms": [1.0] * (1 + len(budgets)), + } + for g, budgets in grid.items() + } + }, + } + ) + ) + spec.update( + max_draft_len=k, + block_size=k, + enable_confidence_scheduling=True, + confidence_cost_table_path=str(grid_path), + confidence_min_predicted_gain=0.0, + ) + config["speculative_config"] = spec + config = without_adp_balance(config) + if args.graph_batch_sizes: + config = with_graph_ladder(config, ladder) + + llm_args = update_llm_args_with_extra_dict({"model": args.model, "backend": "pytorch"}, config) + llm = LLM(**llm_args) + try: + mapping = llm.args.parallel_config.to_mapping() + dp_size = mapping.tp_size if mapping.enable_attention_dp else 1 + max_seq_len = int(llm.args.max_seq_len) + prompts = load_prompts(args.dataset, args.prompt_len) + samples_path = Path(os.environ[PROFILE_OUTPUT_ENV]) + cursor = 0 + capped_tokens = None + for graph_batch_size in sorted(grid, reverse=True): + budgets = [0] + grid[graph_batch_size] + steps = (len(budgets) * args.rounds + 2) * args.hold_steps + for wave in range(args.max_waves): + batch = [] + for _ in range(graph_batch_size * dp_size): + batch.append(prompts[cursor % len(prompts)]) + cursor += 1 + prompt_len = max(len(prompt) for prompt in batch) + max_tokens = min(steps * (k + 1), max_seq_len - prompt_len - 1) + if max_tokens < steps * (k + 1): + # Fine unless acceptance is near-perfect; a wave that drains + # early is followed by another one. + capped_tokens = ( + max_tokens if capped_tokens is None else min(capped_tokens, max_tokens) + ) + if max_tokens < steps: + raise ValueError( + f"max_seq_len={max_seq_len} is too short to hold {steps} decode steps " + f"for G={graph_batch_size}; raise max_seq_len or lower --rounds" + ) + print( + f"[dspark_profile_costs] K={k} G={graph_batch_size} wave {wave + 1}: " + f"{len(batch)} requests x {max_tokens} tokens, V={budgets}", + flush=True, + ) + llm.generate( + [{"prompt_token_ids": prompt} for prompt in batch], + SamplingParams(max_tokens=max_tokens, ignore_eos=True, top_k=1), + use_tqdm=False, + ) + # The runtime schedule resumes where the previous wave stopped. + if not short_cells(samples_path, graph_batch_size, budgets, args.min_samples): + break + setup = CostTableSetup( + gpu=torch.cuda.get_device_name(0), + world_size=int(mapping.world_size), + tensor_parallel_size=int(mapping.tp_size), + moe_expert_parallel_size=int(mapping.moe_ep_size), + enable_attention_dp=bool(mapping.enable_attention_dp), + tensorrt_llm_version=trtllm_version, + ) + finally: + llm.shutdown() + + recorded = json.loads(samples_path.read_text())["samples"] if samples_path.exists() else [] + by_cell = { + (entry["graph_batch_size"], entry["verifier_budget"]): entry["step_time_ms"] + for entry in recorded + } + costs: dict[int, list[tuple[int, float]]] = {} + summary = {} + short = [] + for graph_batch_size, budgets in grid.items(): + cells = [] + for budget in [0] + budgets: + times = by_cell.get((graph_batch_size, budget), []) + if len(times) < args.min_samples: + short.append((graph_batch_size, budget, len(times))) + continue + cells.append((budget, statistics.median(times))) + summary[f"{graph_batch_size}/{budget}"] = len(times) + costs[graph_batch_size] = cells + if short: + hint = "Increase --max-waves or --rounds" + if capped_tokens is not None: + hint = ( + f"max_seq_len capped the waves at {capped_tokens} tokens per request, so they " + "can drain before every tier is timed; raise --max-waves or max_seq_len" + ) + raise RuntimeError( + f"K={k}: too few timed steps for (G, V, samples)={short}; " + f"need {args.min_samples}. {hint}." + ) + payload = build_cost_table_payload( + setup=setup, + costs={k: costs}, + extra={"profile": {str(k): {"samples_per_cell": summary, "hold_steps": args.hold_steps}}}, + ) + Path(args.worker_output).write_text(json.dumps(payload, indent=1)) + for graph_batch_size, cells in costs.items(): + native = dict(cells)[0] + row = ", ".join( + f"V{budget}={ms:.2f}ms ({100.0 * (ms / native - 1.0):+.1f}%)" + for budget, ms in cells + if budget + ) + print( + f"[dspark_profile_costs] K={k} G={graph_batch_size}: native={native:.2f}ms {row}", + flush=True, + ) + + +def _parse_args(argv=None): + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--model", required=True, help="Checkpoint directory.") + parser.add_argument( + "--config", + required=True, + help="Runtime YAML used for serving (parallelism, KV cache, CUDA graphs, DSpark).", + ) + parser.add_argument("--output", required=True, help="Cost table to write or update (JSON).") + parser.add_argument( + "--k", + type=int, + nargs="+", + default=None, + help="Physical draft lengths to profile (default: max_draft_len from --config).", + ) + parser.add_argument( + "--graph-batch-sizes", + type=int, + nargs="+", + default=None, + help="Rank-local CUDA-graph batch sizes to profile; only these graphs are captured " + "(default: the ladder --config captures).", + ) + parser.add_argument( + "--tier-fractions", + type=float, + nargs="+", + default=list(DEFAULT_TIER_FRACTIONS), + help="Compact tiers as the average fraction of K draft tokens verified per request " + "(at least 1/K: every request verifies one draft token).", + ) + parser.add_argument("--hold-steps", type=int, default=16) + parser.add_argument("--warmup-steps", type=int, default=3) + parser.add_argument("--rounds", type=int, default=2) + parser.add_argument("--min-samples", type=int, default=20) + parser.add_argument( + "--max-waves", + type=int, + default=3, + help="Waves per G while a tier has fewer than --min-samples samples.", + ) + parser.add_argument("--prompt-len", type=int, default=64) + parser.add_argument( + "--dataset", + default=None, + help="Optional trtllm-bench JSONL with input_ids to use as prompts.", + ) + parser.add_argument("--workdir", default=None, help="Scratch directory (default: temporary).") + parser.add_argument("--worker-output", default=None, help=argparse.SUPPRESS) + return parser.parse_args(argv) + + +def main(argv=None) -> int: + args = _parse_args(argv) + if args.worker_output: + _profile_k(args) + return 0 + + from tensorrt_llm._torch.speculative.dspark_planner import ( + COST_TABLE_SCHEMA_VERSION, + merge_cost_table_payloads, + ) + + with open(args.config, encoding="utf-8") as file: + config = yaml.safe_load(file) or {} + ks = args.k or [int((config.get("speculative_config") or {}).get("max_draft_len", 0))] + if any(k < 2 for k in ks): + raise SystemExit("--k values must be >= 2 (K1 has no compact verifier budgets)") + ladder = resolve_ladder(config, args.graph_batch_sizes) + for k in ks: + too_small = tiers_below_one_draft_token(ladder, k, args.tier_fractions) + if too_small: + graph_batch_size, fraction, budget = too_small[0] + raise SystemExit( + f"--tier-fractions {fraction:g} gives V={budget} for K={k} at G={graph_batch_size}, " + f"but a full batch verifies at least {2 * graph_batch_size} rows (one draft token " + f"per request); use fractions of at least {1 / k:.2f}" + ) + + output = Path(args.output) + output.parent.mkdir(parents=True, exist_ok=True) + workdir = Path(args.workdir or tempfile.mkdtemp(prefix="dspark-profile-")) + workdir.mkdir(parents=True, exist_ok=True) + merged = json.loads(output.read_text()) if output.exists() else None + if merged is not None and merged.get("schema_version") != COST_TABLE_SCHEMA_VERSION: + raise SystemExit( + f"{output} is not a schema-{COST_TABLE_SCHEMA_VERSION} cost table, so new costs " + "cannot be added to it; choose a new --output path" + ) + for k in ks: + worker_output = workdir / f"costs-k{k}.json" + samples = workdir / f"samples-k{k}.json" + # A rerun in the same --workdir must not count or export an earlier run's samples. + for stale in (worker_output, samples): + stale.unlink(missing_ok=True) + env = dict(os.environ) + env["TLLM_DSPARK_PROFILE_COSTS"] = str(samples) + env["TLLM_DSPARK_PROFILE_HOLD_STEPS"] = str(args.hold_steps) + env["TLLM_DSPARK_PROFILE_WARMUP_STEPS"] = str(args.warmup_steps) + command = [sys.executable, "-m", "tensorrt_llm.tools.dspark_profile_costs"] + command += [arg for arg in (argv if argv is not None else sys.argv[1:]) if arg != ""] + command = _replace_k(command, k) + [ + "--workdir", + str(workdir), + "--worker-output", + str(worker_output), + ] + subprocess.run(command, env=env, check=True) + payload = json.loads(worker_output.read_text()) + merged = payload if merged is None else merge_cost_table_payloads(merged, payload) + tmp_output = output.with_name(output.name + ".tmp") + tmp_output.write_text(json.dumps(merged, indent=1)) + os.replace(tmp_output, output) + print(f"[dspark_profile_costs] wrote K={k} to {output}", flush=True) + return 0 + + +def _replace_k(command: list[str], k: int) -> list[str]: + """Drop any --k/--workdir/--worker-output values from the parent argv and set --k k.""" + result = [] + skip = False + for arg in command: + if skip: + if arg.startswith("--"): + skip = False + else: + continue + if arg in ("--k", "--workdir", "--worker-output"): + skip = True + continue + result.append(arg) + return result + ["--k", str(k)] + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 57cfa122b10b..bd7435c8db6c 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1770,6 +1770,11 @@ "kind": "value", "path": "speculative_config.block_size" }, + { + "capture_policy": "float", + "kind": "value", + "path": "speculative_config.confidence_min_predicted_gain" + }, { "capture_policy": "int|none", "kind": "value", @@ -1831,6 +1836,11 @@ "kind": "categorical", "path": "speculative_config.eagle3_model_arch" }, + { + "capture_policy": "bool", + "kind": "value", + "path": "speculative_config.enable_confidence_scheduling" + }, { "capture_policy": "bool", "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 246e4c0a25e6..0be6dfcf3154 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -140,6 +140,7 @@ l0_cpu: - unittest/_torch/moe - unittest/_torch/multimodal - unittest/_torch/ray_orchestrator/single_gpu/test_cache_transceiver_comm.py + - unittest/_torch/pyexecutor/kv_cache/test_kv_cache_v2_capacity_only.py::test_index_mapper_reservation_tracks_confidence_scheduling - 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. diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_v2_capacity_only.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_v2_capacity_only.py index ec9db1f4067b..55a5e7dd3adc 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_v2_capacity_only.py +++ b/tests/unittest/_torch/pyexecutor/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 ( BlockReusePolicy, @@ -512,3 +513,25 @@ def test_update_of_a_suspended_cache_still_consumes_its_allocation() -> None: assert cache.capacity == 103 assert manager._inflight_draft_lens == {} + + +@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", + ], +) +@pytest.mark.cpu_only +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 diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_v2_vocab_size.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_v2_vocab_size.py index 0b9f3be12053..24d42281235f 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_v2_vocab_size.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_v2_vocab_size.py @@ -21,7 +21,7 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor -from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig, KvCacheConfig from tensorrt_llm.mapping import Mapping pytestmark = pytest.mark.cpu_only @@ -85,7 +85,14 @@ def _dense_config(): def _capture_manager_kwargs( - pretrained_config, base_cls, *, is_multimodal=None, disable_mm_encoder=False + pretrained_config, + base_cls, + *, + is_multimodal=None, + disable_mm_encoder=False, + spec_config=None, + estimating_kv_cache=False, + initialize=False, ): """Route a config through _create_kv_cache_manager, capture the ctor kwargs.""" captured: dict[str, object] = {} @@ -93,6 +100,8 @@ def _capture_manager_kwargs( class RecordingManager(base_cls): def __init__(self, *args: object, **kwargs: object) -> None: captured["kwargs"] = kwargs + if initialize: + super().__init__(*args, **kwargs) model_config = SimpleNamespace( pretrained_config=pretrained_config, @@ -117,7 +126,7 @@ def __init__(self, *args: object, **kwargs: object) -> None: tokens_per_block=64, max_seq_len=2048, max_batch_size=4, - spec_config=None, + spec_config=spec_config, sparse_attention_config=None, max_num_tokens=256, max_beam_width=1, @@ -125,10 +134,81 @@ def __init__(self, *args: object, **kwargs: object) -> None: model_config=model_config, dtype=torch.bfloat16, is_draft=False, + estimating_kv_cache=estimating_kv_cache, ) return captured["kwargs"] +@pytest.mark.parametrize("estimating_kv_cache", [False, True]) +@pytest.mark.parametrize( + "config_factory, base_cls", + [ + pytest.param(_nemotron_hybrid_config, MambaHybridCacheManagerV2, id="nemotron_hybrid"), + pytest.param(_qwen3_5_hybrid_config, MambaHybridCacheManagerV2, id="qwen3_5_hybrid"), + pytest.param(_dense_config, KVCacheManagerV2, id="dense"), + ], +) +def test_factory_preserves_native_index_slot_reservation( + config_factory, base_cls, estimating_kv_cache +): + kwargs = _capture_manager_kwargs( + config_factory(), base_cls, estimating_kv_cache=estimating_kv_cache + ) + + assert "num_reserved_index_slots" not in kwargs + + +@pytest.mark.parametrize( + "enabled, estimating_kv_cache, expected", + [(False, False, None), (False, True, None), (True, False, 2), (True, True, 0)], +) +def test_factory_only_overrides_index_slots_for_enabled_confidence( + enabled, estimating_kv_cache, expected +): + # Validation requires a cost-table path; this test never loads it. + spec_config = DSparkDecodingConfig( + max_draft_len=3, + enable_confidence_scheduling=enabled, + confidence_cost_table_path="costs.json", + ) + kwargs = _capture_manager_kwargs( + _dense_config(), + KVCacheManagerV2, + spec_config=spec_config, + estimating_kv_cache=estimating_kv_cache, + ) + + if expected is None: + assert "num_reserved_index_slots" not in kwargs + else: + assert kwargs["num_reserved_index_slots"] == expected + + +@pytest.mark.parametrize("config_factory", [_nemotron_hybrid_config, _qwen3_5_hybrid_config]) +@pytest.mark.parametrize("estimating_kv_cache", [False, True]) +def test_hybrid_constructor_receives_its_own_index_slot_reservation( + monkeypatch: pytest.MonkeyPatch, config_factory, estimating_kv_cache +): + reservations = [] + + def capture_base_init(manager, *args, num_reserved_index_slots, **kwargs): + reservations.append((num_reserved_index_slots, manager._num_reserved_dummy_slots)) + + monkeypatch.setattr(KVCacheManagerV2, "__init__", capture_base_init) + monkeypatch.setattr(MambaHybridCacheManagerV2, "_setup_state_pool", lambda self: None) + _capture_manager_kwargs( + config_factory(), + MambaHybridCacheManagerV2, + estimating_kv_cache=estimating_kv_cache, + initialize=True, + ) + + assert len(reservations) == 1 + received, owned = reservations[0] + assert received == owned + assert received > 0 + + @pytest.mark.parametrize( "is_multimodal, disable_mm_encoder, expected_vocab_size", [ diff --git a/tests/unittest/_torch/pyexecutor/test_adp_host_exchange.py b/tests/unittest/_torch/pyexecutor/test_adp_host_exchange.py index 050e69714238..2be4ef3a4918 100644 --- a/tests/unittest/_torch/pyexecutor/test_adp_host_exchange.py +++ b/tests/unittest/_torch/pyexecutor/test_adp_host_exchange.py @@ -210,8 +210,12 @@ def test_can_queue_gathers_only_the_batch_size(): """Votes on the batch size alone, over TP x CP; one empty rank vetoes the forward.""" executor = object.__new__(PyExecutor) executor.enable_attention_dp = True + executor.model_engine = SimpleNamespace(_dspark_confidence_enabled=False) executor.dist = SimpleNamespace( - cp_size=1, tp_rank=0, tp_cp_allgather_int64=Mock(return_value=np.array([[2], [0]])) + tp_size=2, + cp_size=1, + tp_rank=0, + tp_cp_allgather_int64=Mock(return_value=np.array([[2], [0]])), ) can_queue, this_rank = PyExecutor._can_queue(executor, types.SimpleNamespace(batch_size=2)) @@ -226,9 +230,13 @@ def test_can_queue_this_rank_is_shared_by_the_cp_peers_of_a_dp_rank(batch_size): vote fails for both, and both defer the previous batch together.""" executor = object.__new__(PyExecutor) executor.enable_attention_dp = True + executor.model_engine = SimpleNamespace(_dspark_confidence_enabled=False) # tp1 x cp2: rows are this DP rank's two CP ranks. executor.dist = SimpleNamespace( - cp_size=2, tp_rank=0, tp_cp_allgather_int64=Mock(return_value=np.array([[1], [0]])) + tp_size=1, + cp_size=2, + tp_rank=0, + tp_cp_allgather_int64=Mock(return_value=np.array([[1], [0]])), ) can_queue, this_rank = PyExecutor._can_queue( diff --git a/tests/unittest/_torch/pyexecutor/test_kv_connector_executor_lifetime.py b/tests/unittest/_torch/pyexecutor/test_kv_connector_executor_lifetime.py index b056d62e2493..6eb32bce051d 100644 --- a/tests/unittest/_torch/pyexecutor/test_kv_connector_executor_lifetime.py +++ b/tests/unittest/_torch/pyexecutor/test_kv_connector_executor_lifetime.py @@ -30,6 +30,7 @@ def _executor() -> PyExecutor: executor = object.__new__(PyExecutor) executor.enable_attention_dp = False + executor.model_engine = SimpleNamespace(_dspark_confidence_enabled=False) executor._kv_connector_failed = False executor._kv_connector_has_pending = False executor._closed = False diff --git a/tests/unittest/_torch/pyexecutor/test_kv_connector_initialization.py b/tests/unittest/_torch/pyexecutor/test_kv_connector_initialization.py index c8c6d0db713a..3b28dacc72ab 100644 --- a/tests/unittest/_torch/pyexecutor/test_kv_connector_initialization.py +++ b/tests/unittest/_torch/pyexecutor/test_kv_connector_initialization.py @@ -86,6 +86,7 @@ def test_connector_retention_request_validation( ) -> None: executor = object.__new__(PyExecutor) executor._is_kv_manager_v2 = use_v2 + executor.model_engine = SimpleNamespace(_dspark_confidence_enabled=False) executor.kv_connector_manager = ( None if connector_mode is None else Mock(capacity_only=connector_mode == "capacity") ) diff --git a/tests/unittest/_torch/pyexecutor/test_py_executor.py b/tests/unittest/_torch/pyexecutor/test_py_executor.py index 7167cccd3bc3..27869ecd8b34 100644 --- a/tests/unittest/_torch/pyexecutor/test_py_executor.py +++ b/tests/unittest/_torch/pyexecutor/test_py_executor.py @@ -2723,6 +2723,8 @@ def _make_adp_request( class _StubADPExecutor: + _dspark_adp_shape_cache_enabled = PyExecutor._dspark_adp_shape_cache_enabled + def __init__( self, *, @@ -2765,7 +2767,11 @@ def __init__( else enable_non_overlap_adp_forward_intent ) self.add_dummy_calls = [] - self.model_engine = Mock(max_num_tokens=max_num_tokens, max_seq_len=max_seq_len) + self.model_engine = Mock( + max_num_tokens=max_num_tokens, + max_seq_len=max_seq_len, + _dspark_confidence_enabled=False, + ) self.dist = Mock() # Peers are modeled through peer_forward_intent; tp_size only tells the diff --git a/tests/unittest/_torch/pyexecutor/test_remote_tail_adp_phase.py b/tests/unittest/_torch/pyexecutor/test_remote_tail_adp_phase.py index 1a27e328133f..5dbd1c39a8b0 100644 --- a/tests/unittest/_torch/pyexecutor/test_remote_tail_adp_phase.py +++ b/tests/unittest/_torch/pyexecutor/test_remote_tail_adp_phase.py @@ -32,11 +32,13 @@ def _make_executor() -> PyExecutor: executor = PyExecutor.__new__(PyExecutor) + executor.enable_attention_dp = True executor.inflight_req_ids = set() executor._remote_tail_phase = RemoteTailPhase.NONE executor._last_remote_tail_phase = RemoteTailPhase.DECODE executor.model_engine = SimpleNamespace( - cuda_graph_runner=SimpleNamespace(padding_dummy_requests={0: object()}) + _dspark_confidence_enabled=False, + cuda_graph_runner=SimpleNamespace(padding_dummy_requests={0: object()}), ) executor.scheduler = Mock() executor.scheduler.is_request_in_schedulable_state.side_effect = ( @@ -346,7 +348,7 @@ def test_queue_requires_execution_rows_and_some_uncancelled_semantic_work( ) -> None: executor = _make_executor() executor._enable_remote_tail_adp = True - executor.dist = Mock() + executor.dist = Mock(tp_size=2) executor.dist.tp_allgather_int64.return_value = torch.tensor(sizes) batch = ScheduledRequests() batch.attention_dp_phase = RemoteTailPhase.DECODE @@ -364,7 +366,7 @@ def test_queue_reports_lost_execution_reservation_collectively( ) -> None: executor = _make_executor() executor._enable_remote_tail_adp = True - executor.dist = Mock() + executor.dist = Mock(tp_size=2) if local_reservation_lost: executor.model_engine.cuda_graph_runner.padding_dummy_requests.clear() executor.dist.tp_allgather_int64.return_value = torch.tensor([[1, 0], [-1, 2]]) @@ -390,6 +392,7 @@ def test_rebalance_reserves_execution_row_before_resuming_real_requests( executor.active_requests = [real] runner = Mock() runner.padding_dummy_requests = {0: dummy} + runner.secondary_padding_dummy_requests = {} executor.model_engine.cuda_graph_runner = runner manager = Mock() executor.kv_cache_manager = manager 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_cost_table.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_table.py index e4a35388b097..796b19d27961 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_table.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_table.py @@ -2,49 +2,68 @@ # SPDX-License-Identifier: Apache-2.0 """DSpark cost-table and scheduling unit tests (hardware-agnostic, CPU).""" -import hashlib import json import pytest import torch +from tensorrt_llm._torch.speculative import dspark_planner from tensorrt_llm._torch.speculative.dspark_planner import ( + CostTableSetup, ExactSpsCostRow, ExactSpsCostTable, - ExactSpsDrainGuard, + build_cost_table_payload, load_runtime_sps_cost_table, + merge_cost_table_payloads, select_exact_sps_candidate, - validate_sps_cost_table_payload, ) -from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig, compute_survival - -BLOCK = 7 - - -def _drain_guard( - *, - tail_graph_batch_size=128, - loss_multiplier=1.0, - mean_output_tokens_per_request_iteration=10.0, - minimum_group_value_ms=0.0, - source_result_sha256="d" * 64, -): - return ExactSpsDrainGuard( - loss_multiplier=loss_multiplier, - mean_output_tokens_per_request_iteration=(mean_output_tokens_per_request_iteration), - minimum_group_value_ms=minimum_group_value_ms, - tail_graph_batch_size=tail_graph_batch_size, - source_result_sha256=source_result_sha256, +from tensorrt_llm._torch.speculative.dspark_schedule import compute_survival + +SETUP = CostTableSetup( + gpu="NVIDIA B300", + world_size=8, + tensor_parallel_size=8, + moe_expert_parallel_size=8, + enable_attention_dp=True, + tensorrt_llm_version="1.4.0", +) + + +def _payload(): + """Schema-3 table with K3 and K5 entries.""" + return build_cost_table_payload( + setup=SETUP, + costs={ + 3: {64: [(0, 44.0), (192, 41.0)], 128: [(0, 60.0), (384, 53.0)]}, + 5: { + 16: [(0, 33.6), (48, 30.8)], + 64: [(0, 49.2), (224, 42.7), (256, 43.6)], + 128: [(0, 75.6), (448, 57.2), (512, 59.7), (576, 65.6), (640, 69.6)], + }, + }, ) -def _cfg(**kw) -> DSparkScheduleConfig: - return DSparkScheduleConfig(block_size=BLOCK, **kw) +def _write(tmp_path, payload, name="costs.json"): + path = tmp_path / name + path.write_text(json.dumps(payload)) + return path -# -------------------------------------------------------------------------- -# survival -# -------------------------------------------------------------------------- +def _load(path, *, k=5, graphs=(64, 128), gain=0.04, live=SETUP): + return load_runtime_sps_cost_table( + path, + graph_batch_sizes=list(graphs), + max_draft_len=k, + minimum_predicted_gain=gain, + live_setup=live, + ) + + +def _with(setup, **changes): + values = setup.to_json() + values.update(changes) + return CostTableSetup(**values) def test_survival_is_cumulative_product(): @@ -53,82 +72,180 @@ def test_survival_is_cumulative_product(): assert torch.allclose(surv[0, :3], torch.tensor([0.9, 0.72, 0.36]), atol=1e-6) -def _multi_g_sps_payload(): - fingerprint = { - "gpu": "B300", - "gpu_count": 8, - "gpu_snapshot_sha256": "a" * 64, - "global_graph_batch_sizes": [512, 1024], - "max_draft_len": 5, - "rank_local_graph_batch_sizes": [64, 128], - "runtime_snapshot": "runtime-v2", - "source_diff_sha256": "b" * 64, - "source_head": "23b73d8", - "topology": "DEP8", - } - cells = { - 64: ((0, 352), (6.0, 5.2)), - 128: ((0, 704, 736), (8.0, 7.0, 7.2)), - } - payload = { +# -------------------------------------------------------------------------- +# Loading: one file for several K, captured-ladder coverage, setup checks +# -------------------------------------------------------------------------- + + +def test_loader_selects_the_runtime_k_and_never_interpolates(tmp_path): + path = _write(tmp_path, _payload()) + + table = _load(path, k=5, gain=0.02) + + assert isinstance(table, ExactSpsCostTable) + assert table.max_draft_len == 5 + assert table.minimum_predicted_gain == pytest.approx(0.02) + assert table.candidate_budgets(128, include_native=True) == (0, 448, 512, 576, 640) + assert table.step_time(448, 128) == pytest.approx(57.2) + with pytest.raises(ValueError, match=r"no direct measurements for G=128, V=\[450\]"): + table.step_time(450, 128) + k3 = _load(path, k=3) + assert k3.max_draft_len == 3 + assert k3.step_time(384, 128) == pytest.approx(53.0) + + +def test_loader_installs_only_captured_graph_sizes(tmp_path): + table = _load(_write(tmp_path, _payload()), graphs=(128,)) + assert sorted(table.tables) == [128] + assert table.candidate_cells() == ((128, 448), (128, 512), (128, 576), (128, 640)) + + +def test_loader_reports_a_missing_k_with_a_profile_command(tmp_path): + path = _write(tmp_path, _payload()) + with pytest.raises( + ValueError, match=r"no costs for max_draft_len=4 \(it has K=\[3, 5\]\)" + ) as error: + _load(path, k=4) + assert "dspark_profile_costs" in str(error.value) + assert "--k 4" in str(error.value) + + +def test_loader_reports_uncovered_captured_graph_sizes(tmp_path): + path = _write(tmp_path, _payload()) + with pytest.raises( + ValueError, match=r"captured graph batch sizes \[32\]; it covers \[16, 64, 128\]" + ): + _load(path, graphs=(32, 64, 128)) + + +def test_loader_reports_a_missing_file(tmp_path): + with pytest.raises(FileNotFoundError, match="does not exist. Generate one with"): + _load(tmp_path / "missing.json") + + +@pytest.mark.parametrize( + "field, value", + [ + ("gpu", "NVIDIA B200"), + ("world_size", 4), + ("tensor_parallel_size", 4), + ("moe_expert_parallel_size", 4), + ("enable_attention_dp", False), + ], +) +def test_loader_rejects_tables_from_another_setup(tmp_path, field, value): + path = _write(tmp_path, _payload()) + with pytest.raises(ValueError, match=f"different setup \\({field}: "): + _load(path, live=_with(SETUP, **{field: value})) + + +def test_loader_warns_on_another_tensorrt_llm_version(tmp_path, monkeypatch): + warnings = [] + monkeypatch.setattr(dspark_planner.logger, "warning", warnings.append) + _load(_write(tmp_path, _payload()), live=_with(SETUP, tensorrt_llm_version="1.5.0")) + assert len(warnings) == 1 + assert "measured with TensorRT-LLM 1.4.0; running 1.5.0" in warnings[0] + + +def _legacy_v2_payload(max_draft_len=5): + return { "schema_version": 2, - "minimum_predicted_gain": 0.02, + "minimum_predicted_gain": 0.04, "cost_tables": { - str(graph_batch_size): { - "token_counts": list(verifier_budgets), - "step_time_ms": list(step_times), - } - for graph_batch_size, (verifier_budgets, step_times) in cells.items() + "64": {"token_counts": [0, 224], "step_time_ms": [49.2, 42.7]}, + "128": {"token_counts": [0, 448], "step_time_ms": [75.6, 57.2]}, + }, + "engine_fingerprint": { + "gpu": "NVIDIA B300", + "gpu_count": 8, + "max_draft_len": max_draft_len, + "rank_local_graph_batch_sizes": [64, 128], + "topology": json.dumps({"attention_dp": True, "cp": 1, "ep": 8, "pp": 1, "tp": 8}), }, - "engine_fingerprint": fingerprint, - "measurements": [ - { - "rank_local_graph_batch_size": graph_batch_size, - "rank_local_verifier_budget": verifier_budget, - "step_time_ms": step_time, - "source_result_sha256": "c" * 64, - } - for graph_batch_size, (verifier_budgets, step_times) in cells.items() - for verifier_budget, step_time in zip(verifier_budgets, step_times) - ], + "engine_fingerprint_sha256": "0" * 64, + "measurements": [], } - payload["engine_fingerprint_sha256"] = hashlib.sha256( - json.dumps(fingerprint, sort_keys=True, separators=(",", ":")).encode() - ).hexdigest() - return payload -def _load_test_exact(path, payload): - fingerprint_path = path.with_name(f"{path.stem}-live-fingerprint.json") - fingerprint_path.write_text(json.dumps(payload["engine_fingerprint"])) - return load_runtime_sps_cost_table( - path, - graph_batch_sizes=[64, 128], - max_draft_len=5, - live_engine_fingerprint_path=fingerprint_path, +def test_loader_reads_legacy_schema_v2_tables(tmp_path): + table = _load(_write(tmp_path, _legacy_v2_payload()), gain=0.0) + assert table.step_time(448, 128) == pytest.approx(57.2) + # The runtime gain comes from the config, not from the legacy file. + assert table.minimum_predicted_gain == 0.0 + + +def test_loader_rejects_a_legacy_table_for_another_k(tmp_path): + path = _write(tmp_path, _legacy_v2_payload(max_draft_len=3)) + with pytest.raises(ValueError, match="measured for max_draft_len=3, but the runtime uses"): + _load(path) + + +def test_loader_checks_legacy_topology(tmp_path): + path = _write(tmp_path, _legacy_v2_payload()) + with pytest.raises(ValueError, match="moe_expert_parallel_size: table=8, runtime=4"): + _load(path, live=_with(SETUP, moe_expert_parallel_size=4)) + + +def test_loader_rejects_unknown_schema(tmp_path): + path = _write(tmp_path, {"schema_version": 7}) + with pytest.raises(ValueError, match="unsupported schema_version=7"): + _load(path) + + +# -------------------------------------------------------------------------- +# Writing and merging tables +# -------------------------------------------------------------------------- + + +def test_build_rejects_cells_the_runtime_would_reject(): + with pytest.raises(ValueError, match=r"G <= V <= G\*\(K\+1\)"): + build_cost_table_payload(setup=SETUP, costs={5: {128: [(0, 75.0), (900, 70.0)]}}) + with pytest.raises(ValueError, match="V=0 native"): + build_cost_table_payload(setup=SETUP, costs={5: {128: [(448, 57.0)]}}) + + +def test_merge_replaces_profiled_rows_and_keeps_the_others(): + existing = _payload() + new = build_cost_table_payload( + setup=SETUP, + costs={5: {128: [(0, 70.0), (448, 55.0)], 256: [(0, 120.0)]}, 4: {128: [(0, 66.0)]}}, ) + merged = merge_cost_table_payloads(existing, new) + assert list(merged["costs"]) == ["3", "4", "5"] + assert merged["costs"]["3"] == existing["costs"]["3"] + assert list(merged["costs"]["5"]) == ["16", "64", "128", "256"] + assert merged["costs"]["5"]["64"] == existing["costs"]["5"]["64"] + assert merged["costs"]["5"]["128"] == {"token_counts": [0, 448], "step_time_ms": [70.0, 55.0]} + + +def test_merge_keeps_per_k_profile_metadata(): + def profiled(k, graph_batch_size, step_ms, samples): + return build_cost_table_payload( + setup=SETUP, + costs={k: {graph_batch_size: [(0, step_ms)]}}, + extra={"profile": {str(k): {"samples_per_cell": {f"{graph_batch_size}/0": samples}}}}, + ) + merged = merge_cost_table_payloads(profiled(4, 128, 66.0, 30), profiled(3, 128, 58.0, 26)) + assert merged["profile"] == { + "4": {"samples_per_cell": {"128/0": 30}}, + "3": {"samples_per_cell": {"128/0": 26}}, + } + merged = merge_cost_table_payloads(merged, profiled(4, 256, 120.0, 12)) + assert merged["profile"]["4"] == {"samples_per_cell": {"128/0": 30, "256/0": 12}} -def test_exact_cost_table_never_interpolates(tmp_path): - payload = _multi_g_sps_payload() - path = tmp_path / "multi-g-sps.json" - path.write_text(json.dumps(payload)) - table, _ = _load_test_exact(path, payload) +def test_merge_refuses_tables_from_different_setups(): + other = build_cost_table_payload( + setup=_with(SETUP, gpu="NVIDIA B200"), costs={5: {128: [(0, 70.0)]}} + ) + with pytest.raises(ValueError, match="different setups: gpu"): + merge_cost_table_payloads(_payload(), other) - assert isinstance(table, ExactSpsCostTable) - assert table.minimum_predicted_gain == pytest.approx(0.02) - assert table.candidate_budgets(64) == (352,) - assert table.candidate_budgets(128) == (704, 736) - assert table.candidate_budgets(128, include_native=True) == (0, 704, 736) - assert table.step_time(0, 128) == pytest.approx(8.0) - assert table.step_time(704, 128) == pytest.approx(7.0) - with pytest.raises(ValueError, match=r"no direct measurements for G=128, V=\[720\]"): - table.step_time(720, 128) - with pytest.raises(ValueError, match="no direct measurements for G=256"): - table.step_time(704, 256) - with pytest.raises(TypeError, match="requested verifier budget"): - table.step_time(704.5, 128) + +# -------------------------------------------------------------------------- +# Table invariants +# -------------------------------------------------------------------------- def test_exact_cost_table_rejects_fractional_programmatic_cells(): @@ -167,12 +284,6 @@ def test_exact_cost_table_identity_covers_grid_costs_and_policy(): minimum_predicted_gain=0.02, ) changed_policy = ExactSpsCostTable(tables=tables, max_draft_len=5, minimum_predicted_gain=0.03) - guarded = ExactSpsCostTable( - tables=tables, - max_draft_len=5, - minimum_predicted_gain=0.02, - iteration_drain_guard=_drain_guard(tail_graph_batch_size=64), - ) assert first.identity_sha256 == reordered.identity_sha256 assert first.collective_identity_words == reordered.collective_identity_words @@ -180,23 +291,6 @@ def test_exact_cost_table_identity_covers_grid_costs_and_policy(): assert all(0 <= value <= 0xFFFFFFFF for value in first.collective_identity_words) assert first.identity_sha256 != changed_cost.identity_sha256 assert first.identity_sha256 != changed_policy.identity_sha256 - assert first.identity_sha256 != guarded.identity_sha256 - - -def test_iteration_drain_guard_rejects_invalid_or_unmeasured_metadata(): - with pytest.raises(ValueError, match="loss_multiplier"): - _drain_guard(loss_multiplier=0.0) - with pytest.raises(ValueError, match="mean_output_tokens"): - _drain_guard(mean_output_tokens_per_request_iteration=float("nan")) - with pytest.raises(ValueError, match="minimum_group_value_ms"): - _drain_guard(minimum_group_value_ms=-0.1) - - with pytest.raises(ValueError, match="measured native tail table for G=16"): - ExactSpsCostTable( - tables={128: ExactSpsCostRow(token_counts=(0,), step_time_ms=(8.0,))}, - max_draft_len=5, - iteration_drain_guard=_drain_guard(tail_graph_batch_size=16), - ) @pytest.mark.parametrize("invalid_budget", [127, 769]) @@ -210,354 +304,76 @@ def test_exact_cost_table_direct_construction_bounds_positive_cells(invalid_budg ) -def test_runtime_loader_rejects_legacy_cost_curve(tmp_path): - path = tmp_path / "legacy.json" - path.write_text( - json.dumps( - { - "token_counts": [0, 100], - "step_time_ms": [1.0, 5.0], - "_meta": {"lookup": "interp"}, - } - ) - ) - - with pytest.raises(ValueError, match="schema_version=2"): - load_runtime_sps_cost_table( - path, - graph_batch_sizes=[64, 128], - max_draft_len=5, - ) - - -def test_exact_cost_table_loader_requires_schema_v2(tmp_path): - payload = _multi_g_sps_payload() - payload["schema_version"] = 1 - path = tmp_path / "wrong-version.json" - path.write_text(json.dumps(payload)) - with pytest.raises(ValueError, match="schema_version=2"): - _load_test_exact(path, payload) - - payload = _multi_g_sps_payload() - del payload["cost_tables"] - path = tmp_path / "missing-tables.json" - path.write_text(json.dumps(payload)) - with pytest.raises(ValueError, match="missing required fields: cost_tables"): - _load_test_exact(path, payload) - payload = _multi_g_sps_payload() - del payload["schema_version"] - path = tmp_path / "missing-version.json" - path.write_text(json.dumps(payload)) - with pytest.raises(ValueError, match="schema_version=2"): - _load_test_exact(path, payload) - - -def test_runtime_exact_loader_requires_independent_fingerprint_file(tmp_path): - payload = _multi_g_sps_payload() - table_path = tmp_path / "exact-sps.json" - table_path.write_text(json.dumps(payload)) - with pytest.raises(ValueError, match="independently generated"): - load_runtime_sps_cost_table( - table_path, - graph_batch_sizes=[64, 128], - max_draft_len=5, - ) - - -@pytest.mark.parametrize("marker", ["engine_fingerprint", "measurements"]) -def test_partial_v2_payload_cannot_bypass_schema_validation(tmp_path, marker): - exact_payload = _multi_g_sps_payload() - mixed_payload = { - "token_counts": [0, 100], - "step_time_ms": [1.0, 5.0], - marker: exact_payload[marker], - } - path = tmp_path / f"mixed-{marker}.json" - path.write_text(json.dumps(mixed_payload)) - fingerprint_path = tmp_path / f"mixed-{marker}-live-fingerprint.json" - fingerprint_path.write_text(json.dumps(exact_payload["engine_fingerprint"])) - - with pytest.raises(ValueError, match="schema_version=2"): - load_runtime_sps_cost_table( - path, - graph_batch_sizes=[64, 128], - max_draft_len=5, - live_engine_fingerprint_path=fingerprint_path, - ) - - -def test_schema_v2_validation_accepts_exact_runtime_grid(): - payload = _multi_g_sps_payload() - - fingerprint = validate_sps_cost_table_payload( - payload, - graph_batch_sizes=[64, 128], +def test_exact_production_candidates_exclude_full_ragged_control(): + table = ExactSpsCostTable( + tables={128: ExactSpsCostRow(token_counts=(0, 704, 768), step_time_ms=(8.0, 7.0, 8.1))}, max_draft_len=5, - live_engine_fingerprint=dict(payload["engine_fingerprint"]), - ) - - assert fingerprint["source_head"] == "23b73d8" - - -def test_schema_v2_loads_authenticated_iteration_drain_metadata(tmp_path): - payload = _multi_g_sps_payload() - payload["iteration_drain_guard"] = { - "loss_multiplier": 1.5, - "mean_output_tokens_per_request_iteration": 4.25, - "minimum_group_value_ms": 2.0, - "source_result_sha256": "d" * 64, - "tail_graph_batch_size": 64, - } - path = tmp_path / "multi-g-sps-with-drain-guard.json" - path.write_text(json.dumps(payload)) - - table, _ = _load_test_exact(path, payload) - - assert table.iteration_drain_guard == _drain_guard( - tail_graph_batch_size=64, - loss_multiplier=1.5, - mean_output_tokens_per_request_iteration=4.25, - minimum_group_value_ms=2.0, ) + assert table.candidate_budgets(128) == (704, 768) + assert table.production_candidate_budgets(128) == (704,) + assert table.candidate_cells() == ((128, 704),) -def test_schema_v2_iteration_drain_metadata_requires_measured_tail_g(tmp_path): - payload = _multi_g_sps_payload() - payload["iteration_drain_guard"] = { - "loss_multiplier": 1.5, - "mean_output_tokens_per_request_iteration": 4.25, - "minimum_group_value_ms": 2.0, - "source_result_sha256": "d" * 64, - "tail_graph_batch_size": 16, - } - path = tmp_path / "multi-g-sps-missing-tail-g.json" - path.write_text(json.dumps(payload)) - - with pytest.raises(ValueError, match="measured native tail table for G=16"): - _load_test_exact(path, payload) - - -def test_schema_v2_validation_requires_the_runtime_graph_grid(): - payload = _multi_g_sps_payload() - with pytest.raises(ValueError, match="must match exactly"): - validate_sps_cost_table_payload( - payload, - graph_batch_sizes=[64], - max_draft_len=5, - live_engine_fingerprint=dict(payload["engine_fingerprint"]), - ) - - -def test_schema_v2_validation_requires_independent_live_fingerprint(): - payload = _multi_g_sps_payload() - with pytest.raises(ValueError, match="independently supplied"): - validate_sps_cost_table_payload( - payload, - graph_batch_sizes=[64, 128], - max_draft_len=5, - ) - with pytest.raises(ValueError, match="independently of the SPS"): - validate_sps_cost_table_payload( - payload, - graph_batch_sizes=[64, 128], - max_draft_len=5, - live_engine_fingerprint=payload["engine_fingerprint"], - ) - - -def test_schema_v2_validation_authenticates_and_matches_fingerprint(): - payload = _multi_g_sps_payload() - live = dict(payload["engine_fingerprint"]) - payload["engine_fingerprint"]["topology"] = "TEP8" - with pytest.raises(ValueError, match="SHA256"): - validate_sps_cost_table_payload( - payload, - graph_batch_sizes=[64, 128], - max_draft_len=5, - live_engine_fingerprint=live, - ) - - payload = _multi_g_sps_payload() - live = dict(payload["engine_fingerprint"]) - live["runtime_snapshot"] = "different" - with pytest.raises(ValueError, match="does not match active runtime"): - validate_sps_cost_table_payload( - payload, - graph_batch_sizes=[64, 128], - max_draft_len=5, - live_engine_fingerprint=live, - ) - - -def test_schema_v2_validation_requires_native_v0_for_every_g(): - payload = _multi_g_sps_payload() - payload["cost_tables"]["128"]["token_counts"].pop(0) - payload["cost_tables"]["128"]["step_time_ms"].pop(0) - payload["measurements"] = [ - item - for item in payload["measurements"] - if not ( - item["rank_local_graph_batch_size"] == 128 and item["rank_local_verifier_budget"] == 0 - ) - ] - - with pytest.raises(ValueError, match="V=0 native static physical-K"): - validate_sps_cost_table_payload( - payload, - graph_batch_sizes=[64, 128], - max_draft_len=5, - live_engine_fingerprint=dict(payload["engine_fingerprint"]), - ) - - -@pytest.mark.parametrize("invalid_budget", [64, 784]) -def test_schema_v2_validation_bounds_positive_verifier_budgets(invalid_budget): - payload = _multi_g_sps_payload() - original_budget = 704 if invalid_budget == 64 else 736 - table_index = 1 if invalid_budget == 64 else 2 - payload["cost_tables"]["128"]["token_counts"][table_index] = invalid_budget - for measurement in payload["measurements"]: - if ( - measurement["rank_local_graph_batch_size"] == 128 - and measurement["rank_local_verifier_budget"] == original_budget - ): - measurement["rank_local_verifier_budget"] = invalid_budget - - with pytest.raises(ValueError, match=r"G <= V <= G\*\(K\+1\)"): - validate_sps_cost_table_payload( - payload, - graph_batch_sizes=[64, 128], +def test_exact_table_limits_compact_graph_cells_per_g(): + with pytest.raises(ValueError, match="compact V cells per G"): + ExactSpsCostTable( + tables={ + 128: ExactSpsCostRow( + token_counts=(0, 128, 256, 384, 512, 640), + step_time_ms=(8.0, 7.9, 7.8, 7.7, 7.6, 7.5), + ) + }, max_draft_len=5, - live_engine_fingerprint=dict(payload["engine_fingerprint"]), ) -def test_schema_v2_validation_rejects_extra_or_mismatched_provenance(): - payload = _multi_g_sps_payload() - payload["cost_tables"]["128"]["token_counts"].append(752) - payload["cost_tables"]["128"]["step_time_ms"].append(7.4) - with pytest.raises(ValueError, match="cells must match exactly"): - validate_sps_cost_table_payload( - payload, - graph_batch_sizes=[64, 128], - max_draft_len=5, - live_engine_fingerprint=dict(payload["engine_fingerprint"]), - ) - - payload = _multi_g_sps_payload() - payload["measurements"][0]["step_time_ms"] += 0.1 - with pytest.raises(ValueError, match="do not match measurement"): - validate_sps_cost_table_payload( - payload, - graph_batch_sizes=[64, 128], +def test_exact_table_limits_total_compact_cells(): + graphs = (16, 24, 32, 40, 48, 56, 64, 72, 80) + with pytest.raises(ValueError, match="limit of 32 compact"): + ExactSpsCostTable( + tables={ + graph: ExactSpsCostRow( + token_counts=(0, 2 * graph, 3 * graph, 4 * graph, 5 * graph), + step_time_ms=(10.0, 6.0, 7.0, 8.0, 9.0), + ) + for graph in graphs + }, max_draft_len=5, - live_engine_fingerprint=dict(payload["engine_fingerprint"]), ) -@pytest.mark.parametrize( - "mutation,message", - [ - (lambda payload: payload.update({"unknown": 1}), "unknown fields"), - (lambda payload: payload.update({"measurements": None}), "null fields"), - (lambda payload: payload["cost_tables"]["128"].update({"unknown": 1}), "unknown fields"), - (lambda payload: payload["engine_fingerprint"].update({"unknown": 1}), "unknown fields"), - ], -) -def test_schema_v2_rejects_null_and_unknown_fields(tmp_path, mutation, message): - payload = _multi_g_sps_payload() - mutation(payload) - path = tmp_path / "invalid-shape.json" - path.write_text(json.dumps(payload)) - with pytest.raises(ValueError, match=message): - _load_test_exact(path, payload) - - -@pytest.mark.parametrize( - "mutation,message", - [ - ( - lambda payload: payload["cost_tables"]["128"]["token_counts"].__setitem__(1, 704.0), - "JSON integer", - ), - ( - lambda payload: payload["measurements"][1].update( - {"rank_local_verifier_budget": 352.5} - ), - "JSON integer", - ), - ( - lambda payload: payload.update({"minimum_predicted_gain": -0.1}), - "non-negative and finite", - ), - ( - lambda payload: payload.update({"minimum_predicted_gain": float("nan")}), - "non-negative and finite", - ), - ( - lambda payload: payload["cost_tables"]["64"]["step_time_ms"].__setitem__(0, 0.0), - "positive and finite", - ), - ( - lambda payload: payload["cost_tables"]["64"]["step_time_ms"].__setitem__( - 0, float("inf") - ), - "positive and finite", - ), - ], -) -def test_schema_v2_rejects_fractional_or_invalid_numbers(tmp_path, mutation, message): - payload = _multi_g_sps_payload() - mutation(payload) - path = tmp_path / "invalid-number.json" - path.write_text(json.dumps(payload)) - with pytest.raises((TypeError, ValueError), match=message): - _load_test_exact(path, payload) - - -def test_schema_v2_rejects_fractional_graph_key(tmp_path): - payload = _multi_g_sps_payload() - payload["cost_tables"]["128.0"] = payload["cost_tables"].pop("128") - path = tmp_path / "fractional-g.json" - path.write_text(json.dumps(payload)) - with pytest.raises(TypeError, match="canonical positive integer string"): - _load_test_exact(path, payload) +# -------------------------------------------------------------------------- +# Selection +# -------------------------------------------------------------------------- -def test_exact_selector_keeps_native_on_tie_or_below_threshold(): - table = ExactSpsCostTable( +def _single_cell_table(gain): + return ExactSpsCostTable( tables={128: ExactSpsCostRow(token_counts=(0, 704), step_time_ms=(10.0, 8.0))}, max_draft_len=5, - minimum_predicted_gain=0.02, - iteration_drain_guard=_drain_guard(), + minimum_predicted_gain=gain, ) - assert ( - select_exact_sps_candidate( - graph_batch_size=128, - native_expected_yield=10.0, - compact_expected_yields={704: 8.0}, - compact_max_yield_losses_per_request={704: 2.0}, - cost_table=table, - ) - == 0 - ) - assert ( - select_exact_sps_candidate( - graph_batch_size=128, - native_expected_yield=10.0, - compact_expected_yields={704: 8.08}, - compact_max_yield_losses_per_request={704: 1.92}, - cost_table=table, + +def test_selector_keeps_native_on_tie_or_below_the_gain_margin(): + table = _single_cell_table(0.02) + # Equal goodput (8 / 8ms == 10 / 10ms) and +1% are both native. + for compact_yield in (8.0, 8.08): + assert ( + select_exact_sps_candidate( + graph_batch_size=128, + native_expected_yield=10.0, + compact_expected_yields={704: compact_yield}, + cost_table=table, + ) + == 0 ) - == 0 - ) + # +3% clears the 2% margin. assert ( select_exact_sps_candidate( graph_batch_size=128, native_expected_yield=10.0, compact_expected_yields={704: 8.24}, - compact_max_yield_losses_per_request={704: 1.76}, cost_table=table, ) == 704 @@ -567,99 +383,23 @@ def test_exact_selector_keeps_native_on_tie_or_below_threshold(): graph_batch_size=128, native_expected_yield=10.0, compact_expected_yields={0: 10.0, 704: 8.24}, - compact_max_yield_losses_per_request={704: 1.76}, - cost_table=table, - ) - - -def test_exact_selector_fails_closed_without_iteration_drain_metadata(): - table = ExactSpsCostTable( - tables={128: ExactSpsCostRow(token_counts=(0, 512), step_time_ms=(100.0, 80.0))}, - max_draft_len=5, - minimum_predicted_gain=0.0, - ) - - # Compact has positive immediate expected goodput, but the table cannot - # price its workload-specific iteration loss. - assert ( - select_exact_sps_candidate( - graph_batch_size=128, - native_expected_yield=10.0, - compact_expected_yields={512: 8.5}, - compact_max_yield_losses_per_request={512: 1.5}, - cost_table=table, - ) - == 0 - ) - - -def test_exact_selector_applies_strict_iteration_drain_group_value(): - table = ExactSpsCostTable( - tables={ - 16: ExactSpsCostRow(token_counts=(0,), step_time_ms=(40.0,)), - 128: ExactSpsCostRow(token_counts=(0, 512), step_time_ms=(100.0, 80.0)), - }, - max_draft_len=5, - minimum_predicted_gain=0.0, - iteration_drain_guard=_drain_guard( - tail_graph_batch_size=16, - loss_multiplier=1.5, - mean_output_tokens_per_request_iteration=5.0, - minimum_group_value_ms=2.0, - ), - ) - - # T(128,0)-T(128,512) - 1.5*(10-8.5)*T(16,0)/5 == 2ms. - # Equality is deliberately native. - assert ( - select_exact_sps_candidate( - graph_batch_size=128, - native_expected_yield=10.0, - compact_expected_yields={512: 8.5}, - compact_max_yield_losses_per_request={512: 1.5}, cost_table=table, ) - == 0 - ) - # A 1.4-token predicted loss has group value 3.2ms and is admitted. - assert ( - select_exact_sps_candidate( - graph_batch_size=128, - native_expected_yield=10.0, - compact_expected_yields={512: 8.6}, - compact_max_yield_losses_per_request={512: 1.4}, - cost_table=table, - ) - == 512 - ) -@pytest.mark.parametrize("graph_batch_size", [16, 32, 64, 128, 192, 256]) -def test_exact_selector_reranks_every_g_by_guarded_group_value(graph_batch_size): +@pytest.mark.parametrize("graph_batch_size", [16, 64, 128, 256]) +def test_selector_picks_the_highest_goodput_tier(graph_batch_size): table = ExactSpsCostTable( tables={ graph_batch_size: ExactSpsCostRow( - token_counts=( - 0, - 3 * graph_batch_size, - 4 * graph_batch_size, - 5 * graph_batch_size, - ), + token_counts=(0, 3 * graph_batch_size, 4 * graph_batch_size, 5 * graph_batch_size), step_time_ms=(100.0, 60.0, 70.0, 80.0), ) }, max_draft_len=5, - minimum_predicted_gain=0.01, - iteration_drain_guard=_drain_guard( - tail_graph_batch_size=graph_batch_size, - mean_output_tokens_per_request_iteration=20.0, - minimum_group_value_ms=20.0, - ), + minimum_predicted_gain=0.04, ) - - # V=4G wins aggregate immediate goodput, but its worst active rank loses - # enough yield per request to miss the E2E guard. V=3G is the best guarded - # tier. V=5G does not clear the aggregate immediate-goodput threshold. + # Goodput: native 0.100, 3G 0.1033, 4G 0.1071, 5G 0.1000 tokens/ms. assert ( select_exact_sps_candidate( graph_batch_size=graph_batch_size, @@ -669,83 +409,27 @@ def test_exact_selector_reranks_every_g_by_guarded_group_value(graph_batch_size) 4 * graph_batch_size: 7.5, 5 * graph_batch_size: 8.0, }, - compact_max_yield_losses_per_request={ - 3 * graph_batch_size: 0.5, - 4 * graph_batch_size: 3.0, - 5 * graph_batch_size: 0.0, - }, cost_table=table, ) - == 3 * graph_batch_size + == 4 * graph_batch_size ) -def test_exact_selector_fails_closed_without_measured_compact_cell(): - graph_batch_size = 32 +def test_selector_keeps_native_without_a_measured_compact_cell(): table = ExactSpsCostTable( tables={ - 16: ExactSpsCostRow(token_counts=(0,), step_time_ms=(40.0,)), 32: ExactSpsCostRow(token_counts=(0,), step_time_ms=(55.0,)), - 64: ExactSpsCostRow(token_counts=(0,), step_time_ms=(70.0,)), 128: ExactSpsCostRow(token_counts=(0, 512), step_time_ms=(100.0, 80.0)), }, max_draft_len=5, minimum_predicted_gain=0.0, - iteration_drain_guard=_drain_guard(tail_graph_batch_size=16), ) - assert ( select_exact_sps_candidate( - graph_batch_size=graph_batch_size, + graph_batch_size=32, native_expected_yield=10.0, compact_expected_yields={}, - compact_max_yield_losses_per_request={}, cost_table=table, ) == 0 ) - - -def test_exact_production_candidates_exclude_full_ragged_control(): - table = ExactSpsCostTable( - tables={128: ExactSpsCostRow(token_counts=(0, 704, 768), step_time_ms=(8.0, 7.0, 8.1))}, - max_draft_len=5, - ) - assert table.candidate_budgets(128) == (704, 768) - assert table.production_candidate_budgets(128) == (704,) - assert table.candidate_cells() == ((128, 704),) - - -def test_exact_table_limits_compact_graph_cells_per_g(): - with pytest.raises(ValueError, match="compact V cells per G"): - ExactSpsCostTable( - tables={ - 128: ExactSpsCostRow( - token_counts=(0, 128, 256, 384, 512, 640), - step_time_ms=(8.0, 7.9, 7.8, 7.7, 7.6, 7.5), - ) - }, - max_draft_len=5, - ) - - -def test_exact_table_accepts_four_compact_cells_for_six_graph_sizes(): - graphs = (16, 32, 64, 128, 192, 256) - table = ExactSpsCostTable( - tables={ - graph: ExactSpsCostRow( - token_counts=(0, int(2.5 * graph), 3 * graph, 4 * graph, 5 * graph), - step_time_ms=(10.0, 6.0, 7.0, 8.0, 9.0), - ) - for graph in graphs - }, - max_draft_len=5, - ) - - assert len(table.candidate_cells()) == 24 - assert table.candidate_cells()[-4:] == ( - (256, 640), - (256, 768), - (256, 1024), - (256, 1280), - ) 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 1f7f739d23ea..c133c5e0e65b 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 @@ -2,7 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 """CPU tests for DSpark ragged CUDA-graph keys and padding state.""" +from contextlib import contextmanager from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch @@ -18,10 +20,89 @@ cuda_graph_dummy_request_id, ) from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import inputs as inputs_module from tensorrt_llm._torch.pyexecutor.sampler.sampler_common import SampleType +from tensorrt_llm._torch.speculative import dspark_planner from tensorrt_llm.mapping import Mapping +@pytest.mark.cpu_only +@pytest.mark.parametrize("enabled", [False, True], ids=["off", "on"]) +def test_decoder_input_preparer_loads_cost_table_for_live_setup(monkeypatch, enabled): + from tensorrt_llm.version import __version__ as trtllm_version + + buffers = SimpleNamespace( + input_ids_cuda=torch.empty(1, dtype=torch.int32, device="cpu"), + gather_ids_cuda=torch.empty(1, dtype=torch.int32, device="cpu"), + ) + spec_config = SimpleNamespace( + enable_confidence_scheduling=enabled, + enable_fused_confidence_scheduler=False, + confidence_cost_table_path="costs.json", + confidence_min_predicted_gain=0.02, + max_draft_len=4, + ) + config = SimpleNamespace( + is_spec_decode=True, + spec_config=spec_config, + cuda_graph_batch_sizes=[1, 2, 4], + enable_attention_dp=True, + ) + mapping = SimpleNamespace(world_size=8, tp_size=8, moe_ep_size=8) + table = Mock() + table.candidate_cells.return_value = ((1, 2), (2, 4)) + table.collective_identity_words = tuple(range(8)) + load_table = Mock(return_value=table) + device_name = Mock(return_value="test-device") + monkeypatch.setattr(dspark_planner, "load_runtime_sps_cost_table", load_table) + monkeypatch.setattr(torch.cuda, "get_device_name", device_name) + monkeypatch.setattr(inputs_module, "prefer_pinned", lambda: False) + monkeypatch.setattr(inputs_module.spec_step_copies, "is_supported", lambda: False) + + preparer = inputs_module.DecoderInputPreparer( + buffers=buffers, + config=config, + mapping=mapping, + dist=None, + model=torch.nn.Identity(), + get_runtime_tokens_per_gen_step=lambda _draft_len: 5, + iter_states={}, + lora=None, + use_beam_search=False, + steady_gen_cache_supported=False, + ) + + assert preparer._buffers is buffers + assert preparer._gather_ids_buffer.tensor is buffers.gather_ids_cuda + if not enabled: + load_table.assert_not_called() + device_name.assert_not_called() + assert preparer._dspark_sps_cost_table is None + assert preparer.previous_verify_lens is None + return + + load_table.assert_called_once_with( + "costs.json", + graph_batch_sizes=[1, 2, 4], + max_draft_len=4, + minimum_predicted_gain=0.02, + live_setup=dspark_planner.CostTableSetup( + gpu="test-device", + world_size=8, + tensor_parallel_size=8, + moe_expert_parallel_size=8, + enable_attention_dp=True, + tensorrt_llm_version=trtllm_version, + ), + ) + device_name.assert_called_once_with(buffers.input_ids_cuda.device) + table.candidate_cells.assert_called_once_with() + assert preparer._dspark_sps_cost_table is table + assert preparer._dspark_exact_candidate_cells == ((1, 2), (2, 4)) + assert preparer._dspark_exact_identity_words == tuple(range(8)) + assert preparer.previous_verify_lens == [] + + class _Batch: def __init__(self, generation_requests): self.generation_requests = list(generation_requests) @@ -271,6 +352,7 @@ def test_reset_padding_dummies_clears_only_runner_owned_state(runner_type): def test_cached_shape_agreement_keeps_sampling_tier_collective(): runner = _runner_stub() batch = _Batch([SimpleNamespace(py_verify_len=3, py_draft_tokens=[0] * 5)] * 4) + runner.set_policy_iteration(9) runner.adp_shape_agreement = _agreement(batch, peer_sizes=(4, 4)) runner.agreed_ragged_bucket = 16 runner.enabled = True @@ -301,9 +383,82 @@ def allgather(payload): def test_cached_shape_agreement_rejects_a_changed_ragged_key(): runner = _runner_stub() batch = _Batch([SimpleNamespace(py_verify_len=3, py_draft_tokens=[0] * 5)] * 4) + runner.set_policy_iteration(9) runner.adp_shape_agreement = _agreement(batch, peer_sizes=(4, 4)) runner.agreed_ragged_bucket = 15 runner.enable_in_graph_sampling = False with pytest.raises(RuntimeError, match="no longer matches"): runner._agree_graph_shape(batch, True, frozenset()) + + +def test_policy_iteration_invalidates_prior_shape_vote(): + runner = _runner_stub() + batch = _Batch([object()] * 3) + runner.set_policy_iteration(9) + agreement = _agreement(batch) + runner.adp_shape_agreement = agreement + assert runner._reusable_padding_agreement(batch) is agreement + runner.set_policy_iteration(9) + assert runner.adp_shape_agreement is agreement + runner.set_policy_iteration(10) + assert runner.adp_shape_agreement is None + # Even a stale producer cannot authenticate its own epoch at the consumer. + runner.adp_shape_agreement = agreement + assert runner._reusable_padding_agreement(batch) is None + + +def test_stale_replay_vote_rejoins_the_shape_collective(): + runner = _runner_stub() + batch = _Batch([SimpleNamespace(py_verify_len=3, py_draft_tokens=[0] * 5)] * 4) + runner.set_policy_iteration(10) + runner.adp_shape_agreement = _agreement(batch, iteration=9, peer_sizes=(4, 4)) + runner.agreed_ragged_bucket = 16 + runner.enabled = True + runner.enable_in_graph_sampling = False + exchanges = [] + + def exchange(payload): + exchanges.append(payload) + return SimpleNamespace(tolist=lambda: [payload, list(payload)]) + + runner.config = SimpleNamespace( + enable_attention_dp=True, + mapping=SimpleNamespace(tp_size=2), + dist=SimpleNamespace(tp_allgather_int64=exchange), + ) + runner._resolve_sample_type = lambda *_args: SampleType.FAST + assert runner._agree_graph_shape(batch, True, frozenset()) == (True, SampleType.FAST) + assert exchanges == [[True, 4, 16, 5, SampleType.FAST.value]] + + +@pytest.mark.parametrize("failure", [None, "enter", "forward"]) +def test_shape_vote_is_consumed_on_normal_and_exceptional_forward_exit(monkeypatch, failure): + runner = _runner_stub() + batch = _Batch([object()] * 3) + runner.set_policy_iteration(9) + agreement = _agreement(batch) + runner.adp_shape_agreement = agreement + + @contextmanager + def parent_pad(self, scheduled_requests, _resource_manager, _runtime_draft_len): + if failure == "enter": + raise RuntimeError("padding failed") + yield scheduled_requests + + monkeypatch.setattr(CUDAGraphRunner, "pad_batch", parent_pad) + + def forward(): + with runner.pad_batch(batch, object(), 5) as padded: + assert padded is batch + assert runner.adp_shape_agreement is agreement + if failure == "forward": + raise RuntimeError("forward failed") + + if failure is None: + forward() + else: + with pytest.raises(RuntimeError, match="failed"): + forward() + assert runner.adp_shape_agreement is None + assert runner._reusable_padding_agreement(batch) is None 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..2251ca2aafb2 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_dummy_slots_cpu.py @@ -0,0 +1,73 @@ +# 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 + 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..e09bb510eb2c --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_embedded_confidence_loading.py @@ -0,0 +1,103 @@ +# 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() + + +@pytest.mark.parametrize("v41", [False, True]) +def test_entirely_absent_confidence_head_remains_supported(monkeypatch, v41): + wrapper, weights = _wrapper(monkeypatch, v41) + 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..e13e425e163e --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_executor_policy.py @@ -0,0 +1,475 @@ +# 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_exact_bucket, +) + +pytestmark = pytest.mark.cpu_only + + +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_yields_across_ranks(): + 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 = _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} + + +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 = _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} + + +def test_exact_yield_decode_rejects_compact_above_native(): + payloads = [_payload(rows=2, native=10_000_000, compact=(11_000_000,))] + + with pytest.raises(RuntimeError, match="exceeds native yield"): + _decode_dspark_exact_expected_yields( + payloads=payloads, + exact_cells=((128, 512),), + graph_batch_size=128, + yield_scale=1_000_000, + ) + + +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) + + +@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, + max_draft_len=5, + ) + executor = PyExecutor.__new__(PyExecutor) + executor.model_engine = SimpleNamespace( + max_draft_len=5, + max_total_draft_tokens=5, + spec_config=spec_config, + cuda_graph_runner=SimpleNamespace( + adp_shape_agreement=None, set_policy_iteration=lambda _iteration: 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.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.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)] + + +def _policy_executor(): + from unittest.mock import Mock + + executor = _draft_len_executor(confidence_enabled=True) + requests = [ + SimpleNamespace( + py_request_id=index + 100, + is_dummy=False, + is_attention_dp_dummy=False, + py_draft_tokens=[1, 2, 3, 4, 5], + py_needs_onehot_draft_probs=False, + py_verify_len=2, + ) + for index in range(3) + ] + batch = _Batch(requests) + batch.can_run_cuda_graph = True + table = SimpleNamespace( + candidate_cells=lambda: ((4, 12),), + collective_identity_words=tuple(range(8)), + ) + local = SimpleNamespace( + num_requests=3, + native_expected_yield=12.0, + compact_expected_yields=(10.0,), + ) + planner = SimpleNamespace( + cost_table=table, + exact_cost_table=table, + max_verify_len=5, + cfg=SimpleNamespace(min_verify_len=1), + prepare_exact_sps_decision=Mock(return_value=local), + allocate_exact_sps_candidate=Mock(return_value=([3, 2, 3], 11, 1)), + device_windows=False, + ) + worker = SimpleNamespace( + verify_planner=planner, + confidence_row_for=lambda request_id: request_id - 100, + staged_confidence_buffer=Mock(return_value=None), + bump_draft_seq=Mock(), + ) + runner = executor.model_engine.cuda_graph_runner + runner.enabled = True + runner.padding_enabled = True + runner.config = SimpleNamespace(batch_size=4, use_mrope=False) + runner.padding_dummy_requests = {5: object()} + runner.secondary_padding_dummy_requests = {5: object()} + runner._round_up_batch_size_with_draft_len = lambda *_args: 4 + runner._round_up_batch_size = lambda _rows: 4 + runner.agreed_ragged_bucket = 99 + runner.ragged_pad_verify_len = 4 + runner.ragged_zero_real_high_rows = 2 + executor.model_engine._get_spec_worker = lambda: worker + executor.model_engine._dspark_trims_submitted_tokens = True + executor.model_engine._dspark_device_budget = 99 + executor.model_engine.fit_ragged_verify_lens = Mock(return_value=12) + executor.dist = SimpleNamespace(tp_size=1) + return executor, batch, worker + + +def _assert_policy_native_cleanup(executor, batch): + runner = executor.model_engine.cuda_graph_runner + assert executor.model_engine._dspark_device_budget is None + assert runner.agreed_ragged_bucket is None + assert runner.ragged_pad_verify_len == 0 + assert runner.ragged_zero_real_high_rows == 0 + assert all(request.py_verify_len is None for request in batch.generation_requests) + + +def test_unqueueable_policy_clears_ragged_state_without_rewriting_draft_buffers(): + executor, batch, worker = _policy_executor() + + def exchange(payload): + peer = list(payload) + peer[-3] = 0 + return SimpleNamespace(tolist=lambda: [payload, peer]) + + executor.dist = SimpleNamespace(tp_size=2, tp_allgather_int64=exchange) + before = [list(request.py_draft_tokens) for request in batch.generation_requests] + executor._handle_dynamic_draft_len(batch) + + _assert_policy_native_cleanup(executor, batch) + assert executor.runtime_draft_len == 5 + assert not executor.model_engine.cuda_graph_runner.adp_shape_agreement.can_queue + assert [request.py_draft_tokens for request in batch.generation_requests] == before + executor.model_engine.fit_ragged_verify_lens.assert_not_called() + worker.staged_confidence_buffer.assert_not_called() + + +def test_no_compact_budget_clears_state_but_stages_the_next_confidence(monkeypatch): + from tensorrt_llm._torch.speculative import dspark_planner + + executor, batch, worker = _policy_executor() + monkeypatch.setattr(dspark_planner, "select_exact_sps_candidate", lambda **_kwargs: 0) + + assert executor._dspark_confidence_draft_len(batch) == 5 + + _assert_policy_native_cleanup(executor, batch) + assert executor.model_engine.cuda_graph_runner.adp_shape_agreement.ragged_bucket is None + executor.model_engine.fit_ragged_verify_lens.assert_not_called() + worker.staged_confidence_buffer.assert_called_once_with() + worker.bump_draft_seq.assert_not_called() + + +@pytest.mark.parametrize("mismatch", ["bucket", "graph_batch_size"]) +def test_policy_rejects_a_fitted_key_different_from_the_agreed_exact_cell(monkeypatch, mismatch): + from tensorrt_llm._torch.pyexecutor import py_executor + from tensorrt_llm._torch.speculative import dspark_planner + + executor, batch, worker = _policy_executor() + monkeypatch.setattr(dspark_planner, "select_exact_sps_candidate", lambda **_kwargs: 12) + if mismatch == "bucket": + executor.model_engine.fit_ragged_verify_lens.return_value = 13 + message = "selected captured verifier bucket" + else: + monkeypatch.setattr(py_executor, "_dspark_exact_common_graph_batch_size", lambda *_args: 8) + message = "disagrees with the cached ADP padding shape" + + with pytest.raises(RuntimeError, match=message): + executor._dspark_confidence_draft_len(batch) + + executor.model_engine.fit_ragged_verify_lens.assert_called_once() + worker.staged_confidence_buffer.assert_not_called() + + +def test_cost_profiler_schedule_replaces_the_policy(monkeypatch): + from unittest.mock import Mock + + from tensorrt_llm._torch.speculative import dspark_planner + + executor, batch, worker = _policy_executor() + planner = worker.verify_planner + planner.exact_cost_table.production_candidate_budgets = lambda _graph_batch_size: (12,) + executor._dspark_cost_profiler = SimpleNamespace(choose=Mock(return_value=12)) + select = Mock(side_effect=AssertionError("the policy must not run while profiling")) + monkeypatch.setattr(dspark_planner, "select_exact_sps_candidate", select) + + executor._dspark_confidence_draft_len(batch) + + executor._dspark_cost_profiler.choose.assert_called_once_with( + iteration=int(executor.iter_counter), + graph_batch_size=4, + budgets=(0, 12), + all_ranks_full=False, + ) + select.assert_not_called() + fit_kwargs = executor.model_engine.fit_ragged_verify_lens.call_args.kwargs + assert fit_kwargs["exact_shape"] == (4, 12, 1) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_profile_costs.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_profile_costs.py new file mode 100644 index 000000000000..25cf5a1b6e85 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_profile_costs.py @@ -0,0 +1,214 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU tests for the DSpark cost profiler schedule and the profiling tool's grid.""" + +import json + +import pytest + +from tensorrt_llm._torch.speculative.dspark_planner import ExactSpsCostRow, ExactSpsCostTable +from tensorrt_llm._torch.speculative.dspark_profile import DSparkCostProfiler +from tensorrt_llm.tools.dspark_profile_costs import ( + _replace_k, + compact_budgets, + main, + plan_grid, + short_cells, + with_graph_ladder, + without_adp_balance, +) + +pytestmark = pytest.mark.cpu_only + + +def test_default_tiers_reproduce_the_measured_k5_grid(): + fractions = (0.5, 0.6, 0.7, 0.8) + assert compact_budgets(128, 5, fractions) == [448, 512, 576, 640] + assert compact_budgets(64, 5, fractions) == [224, 256, 288, 320] + assert compact_budgets(16, 3, fractions) == [40, 45, 50, 54] + # Rounding can collapse tiers at tiny G; they are deduplicated. + assert compact_budgets(1, 5, fractions) == [4, 5] + # V=28 at G16 would leave some request of a full batch with no draft token. + assert compact_budgets(16, 3, (0.25, 0.5)) == [40] + + +def test_tier_fractions_below_one_draft_token_fail_before_profiling(tmp_path, monkeypatch): + config = tmp_path / "runtime.yaml" + config.write_text( + json.dumps( + { + "speculative_config": {"decoding_type": "DSpark", "max_draft_len": 3}, + "cuda_graph_config": {"batch_sizes": [16, 32]}, + } + ) + ) + + def run(*args, **kwargs): + raise AssertionError("profiling must not start") + + monkeypatch.setattr("subprocess.run", run) + argv = ["--model", "m", "--config", str(config), "--output", str(tmp_path / "costs.json")] + with pytest.raises(SystemExit, match="V=28 for K=3 at G=16"): + main(argv + ["--tier-fractions", "0.25", "0.5"]) + + +def test_rerun_in_the_same_workdir_starts_from_fresh_samples(tmp_path, monkeypatch): + config = tmp_path / "runtime.yaml" + config.write_text( + json.dumps( + { + "speculative_config": {"decoding_type": "DSpark", "max_draft_len": 3}, + "cuda_graph_config": {"batch_sizes": [16]}, + } + ) + ) + workdir = tmp_path / "work" + workdir.mkdir() + (workdir / "samples-k3.json").write_text(json.dumps({"samples": []})) + (workdir / "costs-k3.json").write_text("{}") + payload = {"schema_version": 3, "setup": {}, "costs": {"3": {}}} + + def run(command, env, check): + assert not (workdir / "samples-k3.json").exists() + assert env["TLLM_DSPARK_PROFILE_COSTS"] == str(workdir / "samples-k3.json") + worker_output = command[command.index("--worker-output") + 1] + assert not (workdir / "costs-k3.json").exists() + with open(worker_output, "w") as file: + json.dump(payload, file) + + monkeypatch.setattr("subprocess.run", run) + output = tmp_path / "costs.json" + argv = ["--model", "m", "--config", str(config), "--output", str(output)] + assert main(argv + ["--workdir", str(workdir)]) == 0 + assert json.loads(output.read_text()) == payload + + +def test_grid_gives_large_graphs_priority_within_the_runtime_cell_limit(): + ladder = [1, 2, 4] + [8 * i for i in range(1, 17)] + grid = plan_grid(ladder, 5, (0.5, 0.6, 0.7, 0.8)) + assert sorted(grid) == ladder + assert sum(len(budgets) for budgets in grid.values()) == 32 + assert all(len(grid[g]) == 4 for g in (72, 80, 88, 96, 104, 112, 120, 128)) + assert grid[8] == [] + # The planned grid is a table the runtime accepts. + ExactSpsCostTable( + tables={ + g: ExactSpsCostRow(token_counts=[0] + budgets, step_time_ms=[1.0] * (1 + len(budgets))) + for g, budgets in grid.items() + }, + max_draft_len=5, + ) + + +def test_worker_command_profiles_one_k(): + command = ["python", "-m", "tool", "--model", "m", "--k", "3", "5", "--output", "o.json"] + assert _replace_k(command, 5) == [ + "python", + "-m", + "tool", + "--model", + "m", + "--output", + "o.json", + "--k", + "5", + ] + + +def test_profiler_round_robins_budgets_and_times_full_steps(tmp_path, monkeypatch): + clock = iter(float(step) * 0.05 for step in range(1000)) + monkeypatch.setattr("time.perf_counter", lambda: next(clock)) + output = tmp_path / "samples.json" + profiler = DSparkCostProfiler(str(output), write_samples=True, hold_steps=4, warmup_steps=1) + + chosen = [ + profiler.choose(iteration=i, graph_batch_size=8, budgets=(0, 20, 30), all_ranks_full=True) + for i in range(24) + ] + profiler.flush() + + assert chosen == [0] * 4 + [20] * 4 + [30] * 4 + [0] * 4 + [20] * 4 + [30] * 4 + samples = { + (entry["graph_batch_size"], entry["verifier_budget"]): entry["step_time_ms"] + for entry in json.loads(output.read_text())["samples"] + } + # Two holds per budget, each dropping its first step; the last step of the + # final hold has no successor yet. + assert len(samples[(8, 0)]) == 6 + assert len(samples[(8, 20)]) == 6 + assert len(samples[(8, 30)]) == 5 + assert all(value == pytest.approx(50.0) for values in samples.values() for value in values) + + +def test_profiler_skips_partial_steps_and_iteration_gaps(tmp_path, monkeypatch): + clock = iter(float(step) for step in range(1000)) + monkeypatch.setattr("time.perf_counter", lambda: next(clock)) + profiler = DSparkCostProfiler( + str(tmp_path / "samples.json"), write_samples=False, hold_steps=4, warmup_steps=0 + ) + + assert ( + profiler.choose(iteration=0, graph_batch_size=8, budgets=(0, 20), all_ranks_full=True) == 0 + ) + # A rank without G real requests runs native and is not timed. + assert ( + profiler.choose(iteration=1, graph_batch_size=8, budgets=(0, 20), all_ranks_full=False) == 0 + ) + assert ( + profiler.choose(iteration=2, graph_batch_size=8, budgets=(0, 20), all_ranks_full=True) == 0 + ) + # Iteration 3 had no decode decision (for example a mixed step). + assert ( + profiler.choose(iteration=4, graph_batch_size=8, budgets=(0, 20), all_ranks_full=True) == 0 + ) + + assert dict(profiler._samples) == {(8, 0): [1000.0]} + + +def test_short_cells_reports_budgets_below_the_sample_minimum(tmp_path): + samples = tmp_path / "samples.json" + assert short_cells(samples, 64, [0, 224], 20) == {0: 0, 224: 0} + samples.write_text( + json.dumps( + { + "samples": [ + {"graph_batch_size": 64, "verifier_budget": 0, "step_time_ms": [1.0] * 26}, + {"graph_batch_size": 64, "verifier_budget": 224, "step_time_ms": [1.0] * 17}, + {"graph_batch_size": 128, "verifier_budget": 448, "step_time_ms": [1.0]}, + ] + } + ) + ) + assert short_cells(samples, 64, [0, 224, 256], 20) == {224: 17, 256: 0} + assert short_cells(samples, 64, [0], 20) == {} + + +def test_profiling_turns_off_attention_dp_balancing_only(): + balanced = { + "max_batch_size": 64, + "attention_dp_config": {"enable_balance": True, "batching_wait_iters": 10}, + } + profiled = without_adp_balance(balanced) + assert profiled["attention_dp_config"] == {"enable_balance": False, "batching_wait_iters": 10} + assert profiled["max_batch_size"] == 64 + assert balanced["attention_dp_config"]["enable_balance"] is True + assert without_adp_balance({"max_batch_size": 64}) == {"max_batch_size": 64} + + +def test_graph_batch_size_override_captures_only_the_profiled_graphs(): + config = { + "max_batch_size": 64, + "cuda_graph_config": {"batch_sizes": [16, 32, 64], "enable_padding": True}, + } + profiled = with_graph_ladder(config, [64]) + assert profiled["cuda_graph_config"] == { + "batch_sizes": [64], + "max_batch_size": 64, + "enable_padding": True, + } + assert profiled["max_batch_size"] == 64 + assert config["cuda_graph_config"]["batch_sizes"] == [16, 32, 64] + assert with_graph_ladder({}, [16, 32])["cuda_graph_config"] == { + "batch_sizes": [16, 32], + "max_batch_size": 32, + } 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..409c02eccbe2 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sampling_admission.py @@ -0,0 +1,157 @@ +# 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.metadata import ( + supports_advanced_draft_sampling, + sync_group_all_greedy_sample, + validate_dspark_draft_sampling, +) +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 = SimpleNamespace() + 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 + sync_group_all_greedy_sample( + metadata, runner.mapping, runner.dist, confidence_enabled=enabled + ) + if enabled and rejection and not all(flags): + with pytest.raises(ValueError, match="advanced rejection draft sampling"): + validate_dspark_draft_sampling(metadata, confidence_enabled=enabled) + errors.append(True) + else: + validate_dspark_draft_sampling(metadata, confidence_enabled=enabled) + 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 = SimpleNamespace() + runner._dspark_confidence_enabled = enabled + runner._config = SimpleNamespace( + spec_config=SimpleNamespace( + use_rejection_sampling=rejection, enable_confidence_scheduling=enabled + ) + ) + assert supports_advanced_draft_sampling(runner._config.spec_config) 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 = SimpleNamespace() + 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 + + sync_group_all_greedy_sample(metadata, runner.mapping, runner.dist, confidence_enabled=enabled) + + 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_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 6a978a5aeb08..98b4f58ba560 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -102,9 +102,8 @@ def _make_worker(enable_confidence_scheduling=False): max_draft_len=5, spec_dec_mode=SpeculativeDecodingMode.DSPARK, enable_confidence_scheduling=enable_confidence_scheduling, - enable_fused_confidence_scheduler=False, - confidence_sps_table_path=None, - confidence_sts_path=None, + confidence_cost_table_path=None, + confidence_calibration_path=None, ) from tensorrt_llm.mapping import Mapping @@ -422,17 +421,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_confidence_state.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_confidence_state.py index d94e995df727..c75b4e3f96b3 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_confidence_state.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_confidence_state.py @@ -220,8 +220,8 @@ def test_configured_sts_is_loaded_on_the_resolved_head_and_planner(tmp_path): worker = _worker(1) path = tmp_path / "sts.json" path.write_text(json.dumps({"sts_temperatures": [0.5, 1.0, 2.0]}), encoding="utf-8") - worker.spec_config.confidence_sts_path = str(path) - worker.spec_config.confidence_sps_table_path = None + worker.spec_config.confidence_calibration_path = str(path) + worker.spec_config.confidence_cost_table_path = None worker.spec_config.enable_fused_confidence_scheduler = False head = DSparkConfidenceHead(hidden_size=2, block_size=3) model = SimpleNamespace( @@ -244,18 +244,53 @@ def test_configured_sts_is_loaded_on_the_resolved_head_and_planner(tmp_path): def test_configured_sts_without_a_confidence_head_is_rejected_before_loading(tmp_path): worker = _worker(1) - worker.spec_config.confidence_sts_path = str(tmp_path / "must-not-be-opened.json") - worker.spec_config.confidence_sps_table_path = None + worker.spec_config.confidence_calibration_path = str(tmp_path / "must-not-be-opened.json") + worker.spec_config.confidence_cost_table_path = None worker.spec_config.enable_fused_confidence_scheduler = False with pytest.raises( - ValueError, match="confidence_sts_path is set but no confidence head was found" + ValueError, match="confidence_calibration_path is set but no confidence head was found" ): worker._build_verify_planner(SimpleNamespace(), block_size=3) assert not hasattr(worker, "verify_planner") +def _model_with_head(block_size): + head = DSparkConfidenceHead(hidden_size=2, block_size=block_size) + model = SimpleNamespace( + dspark_model=SimpleNamespace(mtp_layers=[SimpleNamespace(confidence_head=head)]) + ) + return head, model + + +def test_calibration_fitted_at_a_larger_k_serves_a_smaller_k(tmp_path): + worker = _worker(1) + path = tmp_path / "calibration.json" + path.write_text(json.dumps({"sts_temperatures": [0.5, 1.0, 2.0, 3.0, 4.0]}), encoding="utf-8") + worker.spec_config.confidence_calibration_path = str(path) + worker.spec_config.confidence_cost_table_path = None + worker.spec_config.enable_fused_confidence_scheduler = False + head, model = _model_with_head(block_size=3) + + worker._build_verify_planner(model, block_size=3) + + assert torch.equal(head.sts_temperatures, torch.tensor([0.5, 1.0, 2.0])) + + +def test_calibration_shorter_than_k_is_rejected(tmp_path): + worker = _worker(1) + path = tmp_path / "calibration.json" + path.write_text(json.dumps({"sts_temperatures": [0.5, 1.0]}), encoding="utf-8") + worker.spec_config.confidence_calibration_path = str(path) + worker.spec_config.confidence_cost_table_path = None + worker.spec_config.enable_fused_confidence_scheduler = False + _, model = _model_with_head(block_size=3) + + with pytest.raises(ValueError, match="has 2 temperatures but the confidence head scores 3"): + worker._build_verify_planner(model, block_size=3) + + @pytest.mark.parametrize("bias", [False, True]) def test_embedded_checkpoint_load_preserves_optional_confidence_bias(monkeypatch, bias): draft = modeling_dspark.DSv4DSparkDraftModel.__new__(modeling_dspark.DSv4DSparkDraftModel) 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/hw_agnostic/test_dspark_workload_admission.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_workload_admission.py deleted file mode 100644 index f8026e353511..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_workload_admission.py +++ /dev/null @@ -1,406 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""CPU-only tests for DSpark workload-level confidence admission.""" - -import hashlib -import json -from copy import deepcopy -from pathlib import Path - -import pytest - -from tensorrt_llm._torch.speculative.dspark_planner import ( - EXACT_SPS_SELECTOR_IDENTITY_SHA256, - evaluate_confidence_workload_admission, - load_confidence_workload_admission, -) - - -def _multi_g_sps_payload(): - fingerprint = { - "gpu": "B300", - "gpu_count": 8, - "gpu_snapshot_sha256": "a" * 64, - "global_graph_batch_sizes": [512, 1024], - "max_draft_len": 5, - "rank_local_graph_batch_sizes": [64, 128], - "runtime_snapshot": "runtime-v2", - "source_diff_sha256": "b" * 64, - "source_head": "23b73d8", - "topology": "DEP8", - } - cells = { - 64: ((0, 352), (6.0, 5.2)), - 128: ((0, 704, 736), (8.0, 7.0, 7.2)), - } - payload = { - "schema_version": 2, - "minimum_predicted_gain": 0.02, - "cost_tables": { - str(graph_batch_size): { - "token_counts": list(verifier_budgets), - "step_time_ms": list(step_times), - } - for graph_batch_size, (verifier_budgets, step_times) in cells.items() - }, - "engine_fingerprint": fingerprint, - "measurements": [ - { - "rank_local_graph_batch_size": graph_batch_size, - "rank_local_verifier_budget": verifier_budget, - "step_time_ms": step_time, - "source_result_sha256": "c" * 64, - } - for graph_batch_size, (verifier_budgets, step_times) in cells.items() - for verifier_budget, step_time in zip(verifier_budgets, step_times) - ], - } - payload["engine_fingerprint_sha256"] = hashlib.sha256( - json.dumps(fingerprint, sort_keys=True, separators=(",", ":")).encode() - ).hexdigest() - return payload - - -def _write_workload_admission( - tmp_path, - *, - admitted=True, - policy_steps=100, - compact_choices=50, - gross_compact_value_ms_lower_bound=160.0, - fixed_confidence_path_overhead_ms_upper_bound=1.0, - safety_margin_ms=10.0, -): - sps_payload = _multi_g_sps_payload() - sps_path = tmp_path / "sps.json" - sps_path.write_text(json.dumps(sps_payload)) - fingerprint_path = tmp_path / "live-fingerprint.json" - fingerprint_path.write_text(json.dumps(sps_payload["engine_fingerprint"])) - admission = { - "admitted": admitted, - "calibration_result_sha256": "d" * 64, - "compact_choices": compact_choices, - "fixed_confidence_path_overhead_ms_upper_bound": ( - fixed_confidence_path_overhead_ms_upper_bound - ), - "gross_compact_value_ms_lower_bound": gross_compact_value_ms_lower_bound, - "live_engine_fingerprint_sha256": hashlib.sha256(fingerprint_path.read_bytes()).hexdigest(), - "physical_k": 5, - "policy_steps": policy_steps, - "runtime_snapshot": sps_payload["engine_fingerprint"]["runtime_snapshot"], - "safety_margin_ms": safety_margin_ms, - "selector_identity_sha256": EXACT_SPS_SELECTOR_IDENTITY_SHA256, - "selector_replay_sha256": "e" * 64, - "source_diff_sha256": sps_payload["engine_fingerprint"]["source_diff_sha256"], - "source_head": sps_payload["engine_fingerprint"]["source_head"], - "sps_cost_table_sha256": hashlib.sha256(sps_path.read_bytes()).hexdigest(), - "workload_identity_sha256": "f" * 64, - } - receipt = { - "schema_version": 1, - "admission": admission, - "admission_sha256": hashlib.sha256( - json.dumps(admission, sort_keys=True, separators=(",", ":")).encode() - ).hexdigest(), - } - receipt_path = tmp_path / "admission.json" - receipt_path.write_text(json.dumps(receipt)) - return ( - receipt_path, - sps_path, - fingerprint_path, - receipt, - hashlib.sha256(receipt_path.read_bytes()).hexdigest(), - ) - - -def test_workload_admission_has_no_policy_steps_per_choice_fixed_point(): - # A selector that charges H*(N/C) back to each candidate has no stable - # finite C for this profitable workload. Starting from all 50 positive - # marginal choices gives R=2 and keeps only the ten 8-ms choices; R then - # becomes 10 and rejects those too. Aggregate accounting admits because - # 10*8 + 40*1.9 - 100*1 = 56 ms before the explicit margin. - marginal_values_ms = [8.0] * 10 + [1.9] * 40 + [0.0] * 50 - - assert sum(value > 2.0 for value in marginal_values_ms) == 10 - assert sum(value > 10.0 for value in marginal_values_ms) == 0 - admitted, net_value_ms = evaluate_confidence_workload_admission( - policy_steps=100, - compact_choices=50, - gross_compact_value_ms_lower_bound=sum(marginal_values_ms), - fixed_confidence_path_overhead_ms_upper_bound=1.0, - safety_margin_ms=0.0, - ) - - assert admitted is True - assert net_value_ms == pytest.approx(56.0) - - -def test_workload_admission_rejects_stale_ratio_undercharge(): - # A stale R=2 would admit each 8-ms opportunity. If that filtering leaves - # only ten compact choices, however, the workload still pays 100 policy - # steps of overhead: 80 - 100 = -20 ms. The aggregate gate cannot be fooled - # by the stale ratio because R is not an input. - assert 8.0 > 1.0 * 2.0 - admitted, net_value_ms = evaluate_confidence_workload_admission( - policy_steps=100, - compact_choices=10, - gross_compact_value_ms_lower_bound=80.0, - fixed_confidence_path_overhead_ms_upper_bound=1.0, - safety_margin_ms=0.0, - ) - - assert admitted is False - assert net_value_ms == pytest.approx(-20.0) - - -def test_workload_admission_requires_a_compact_choice(): - admitted, net_value_ms = evaluate_confidence_workload_admission( - policy_steps=100, - compact_choices=0, - gross_compact_value_ms_lower_bound=200.0, - fixed_confidence_path_overhead_ms_upper_bound=1.0, - safety_margin_ms=10.0, - ) - - assert admitted is False - assert net_value_ms == pytest.approx(90.0) - - -@pytest.mark.parametrize("gross_value,expected_net", [(110.0, 0.0), (100.0, -10.0)]) -def test_workload_admission_declines_nonpositive_net_value( - tmp_path: Path, gross_value: float, expected_net: float -) -> None: - receipt_path, sps_path, fingerprint_path, _, receipt_sha256 = _write_workload_admission( - tmp_path, admitted=False, gross_compact_value_ms_lower_bound=gross_value - ) - - admission = load_confidence_workload_admission( - receipt_path, - expected_receipt_sha256=receipt_sha256, - sps_cost_table_path=sps_path, - live_engine_fingerprint_path=fingerprint_path, - physical_k=5, - ) - - assert admission.admitted is False - assert admission.net_value_ms_lower_bound == expected_net - - -@pytest.mark.parametrize( - "field,value,error", - [ - ("policy_steps", True, TypeError), - ("policy_steps", 0, ValueError), - ("compact_choices", 101, ValueError), - ("gross_compact_value_ms_lower_bound", float("nan"), ValueError), - ("fixed_confidence_path_overhead_ms_upper_bound", float("inf"), ValueError), - ("safety_margin_ms", -1.0, ValueError), - ], -) -def test_workload_admission_rejects_invalid_economics( - field: str, value: object, error: type[Exception] -) -> None: - inputs = dict( - policy_steps=100, - compact_choices=50, - gross_compact_value_ms_lower_bound=160.0, - fixed_confidence_path_overhead_ms_upper_bound=1.0, - safety_margin_ms=10.0, - ) - inputs[field] = value - - with pytest.raises(error): - evaluate_confidence_workload_admission(**inputs) - - -@pytest.mark.parametrize("artifact", ["sps", "fingerprint"]) -def test_workload_admission_rejects_changed_artifact_bytes(tmp_path: Path, artifact: str) -> None: - receipt_path, sps_path, fingerprint_path, _, receipt_sha256 = _write_workload_admission( - tmp_path - ) - changed = sps_path if artifact == "sps" else fingerprint_path - # Equivalent JSON is still a different artifact from the pinned bytes. - changed.write_bytes(changed.read_bytes() + b"\n") - - with pytest.raises(ValueError, match="SHA256 does not match"): - load_confidence_workload_admission( - receipt_path, - expected_receipt_sha256=receipt_sha256, - sps_cost_table_path=sps_path, - live_engine_fingerprint_path=fingerprint_path, - physical_k=5, - ) - - -def test_workload_admission_parses_the_pinned_receipt_bytes(tmp_path: Path, monkeypatch) -> None: - receipt_path, sps_path, fingerprint_path, receipt, receipt_sha256 = _write_workload_admission( - tmp_path - ) - replacement = deepcopy(receipt) - replacement["admission"]["gross_compact_value_ms_lower_bound"] = 170.0 - replacement["admission_sha256"] = hashlib.sha256( - json.dumps(replacement["admission"], sort_keys=True, separators=(",", ":")).encode() - ).hexdigest() - replacement_bytes = json.dumps(replacement).encode() - original_open = Path.open - replaced = False - - class ReplaceOnClose: - def __init__(self, file): - self.file = file - - def __enter__(self): - return self.file.__enter__() - - def __exit__(self, *args): - result = self.file.__exit__(*args) - with original_open(receipt_path, "wb") as file: - file.write(replacement_bytes) - return result - - def open_with_replacement(path, mode="r", *args, **kwargs): - nonlocal replaced - file = original_open(path, mode, *args, **kwargs) - if path == receipt_path and mode == "rb" and not replaced: - replaced = True - return ReplaceOnClose(file) - return file - - monkeypatch.setattr(Path, "open", open_with_replacement) - admission = load_confidence_workload_admission( - receipt_path, - expected_receipt_sha256=receipt_sha256, - sps_cost_table_path=sps_path, - live_engine_fingerprint_path=fingerprint_path, - physical_k=5, - ) - - assert replaced is True - assert receipt_path.read_bytes() == replacement_bytes - assert admission.gross_compact_value_ms_lower_bound == 160.0 - - -def test_workload_admission_rejects_cost_table_over_runtime_cell_limit(tmp_path: Path) -> None: - receipt_path, sps_path, fingerprint_path, receipt, _ = _write_workload_admission(tmp_path) - sps_payload = json.loads(sps_path.read_text()) - budgets = [0, 64, 128, 192, 256, 320, 352] - times = [6.0, 5.9, 5.8, 5.7, 5.6, 5.5, 5.2] - sps_payload["cost_tables"]["64"] = {"token_counts": budgets, "step_time_ms": times} - sps_payload["measurements"] = [ - { - "rank_local_graph_batch_size": int(graph_batch_size), - "rank_local_verifier_budget": budget, - "step_time_ms": step_time, - "source_result_sha256": "c" * 64, - } - for graph_batch_size, cells in sps_payload["cost_tables"].items() - for budget, step_time in zip(cells["token_counts"], cells["step_time_ms"]) - ] - sps_path.write_text(json.dumps(sps_payload)) - receipt["admission"]["sps_cost_table_sha256"] = hashlib.sha256( - sps_path.read_bytes() - ).hexdigest() - receipt["admission_sha256"] = hashlib.sha256( - json.dumps(receipt["admission"], sort_keys=True, separators=(",", ":")).encode() - ).hexdigest() - receipt_path.write_text(json.dumps(receipt)) - - with pytest.raises(ValueError, match="production limit.*compact V cells per G"): - load_confidence_workload_admission( - receipt_path, - expected_receipt_sha256=hashlib.sha256(receipt_path.read_bytes()).hexdigest(), - sps_cost_table_path=sps_path, - live_engine_fingerprint_path=fingerprint_path, - physical_k=5, - ) - - -@pytest.mark.parametrize("field", ["schema_version", "admitted"]) -def test_workload_admission_rejects_duplicate_json_keys(tmp_path: Path, field: str) -> None: - receipt_path, sps_path, fingerprint_path, _, _ = _write_workload_admission(tmp_path) - text = receipt_path.read_text() - literal = '"schema_version": 1' if field == "schema_version" else '"admitted": true' - assert text.count(literal) == 1 - receipt_path.write_text(text.replace(literal, literal + ", " + literal)) - - with pytest.raises(ValueError, match="duplicate field"): - load_confidence_workload_admission( - receipt_path, - expected_receipt_sha256=hashlib.sha256(receipt_path.read_bytes()).hexdigest(), - sps_cost_table_path=sps_path, - live_engine_fingerprint_path=fingerprint_path, - physical_k=5, - ) - - -def test_workload_admission_authenticates_and_recomputes_positive_decision(tmp_path): - receipt_path, sps_path, fingerprint_path, receipt, receipt_sha256 = _write_workload_admission( - tmp_path - ) - - admission = load_confidence_workload_admission( - receipt_path, - expected_receipt_sha256=receipt_sha256, - sps_cost_table_path=sps_path, - live_engine_fingerprint_path=fingerprint_path, - physical_k=5, - ) - - assert admission.admitted is True - assert admission.net_value_ms_lower_bound == pytest.approx(50.0) - assert admission.admission_identity_sha256 == receipt["admission_sha256"] - assert admission.receipt_sha256 == receipt_sha256 - - -def test_workload_admission_rejects_unpinned_receipt_bytes(tmp_path): - receipt_path, sps_path, fingerprint_path, _, _ = _write_workload_admission(tmp_path) - - with pytest.raises(ValueError, match="pinned SHA256"): - load_confidence_workload_admission( - receipt_path, - expected_receipt_sha256="0" * 64, - sps_cost_table_path=sps_path, - live_engine_fingerprint_path=fingerprint_path, - physical_k=5, - ) - - -@pytest.mark.parametrize( - "mutation,message", - [ - ( - lambda receipt: receipt["admission"].update({"selector_identity_sha256": "0" * 64}), - "selector identity", - ), - ( - lambda receipt: receipt["admission"].update({"physical_k": 4}), - "physical_k", - ), - ( - lambda receipt: receipt["admission"].update({"source_head": "stale"}), - "source_head", - ), - ( - lambda receipt: receipt["admission"].update({"policy_steps": 1000}), - "decision does not match", - ), - ], -) -def test_workload_admission_rejects_stale_or_inconsistent_receipt(tmp_path, mutation, message): - receipt_path, sps_path, fingerprint_path, receipt, _ = _write_workload_admission(tmp_path) - mutation(receipt) - receipt["admission_sha256"] = hashlib.sha256( - json.dumps(receipt["admission"], sort_keys=True, separators=(",", ":")).encode() - ).hexdigest() - receipt_path.write_text(json.dumps(receipt)) - receipt_sha256 = hashlib.sha256(receipt_path.read_bytes()).hexdigest() - - with pytest.raises(ValueError, match=message): - load_confidence_workload_admission( - receipt_path, - expected_receipt_sha256=receipt_sha256, - sps_cost_table_path=sps_path, - live_engine_fingerprint_path=fingerprint_path, - physical_k=5, - ) diff --git a/tests/unittest/disaggregated/test_chunked_transfer.py b/tests/unittest/disaggregated/test_chunked_transfer.py index 61197e997ace..3dff00727583 100644 --- a/tests/unittest/disaggregated/test_chunked_transfer.py +++ b/tests/unittest/disaggregated/test_chunked_transfer.py @@ -614,6 +614,7 @@ def test_pipelined_transfer_requires_gen_first_flow(): from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor executor = MagicMock(_is_kv_manager_v2=True) + executor.model_engine._dspark_confidence_enabled = False executor.is_warmup = False executor.kv_cache_transceiver.pipeline_transfer_enabled = True executor.model_engine.attn_runtime_features.chunked_prefill = True @@ -664,6 +665,7 @@ def test_pipelined_transfer_rejects_pipeline_parallelism_for_context_request(): from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor executor = MagicMock(_is_kv_manager_v2=True) + executor.model_engine._dspark_confidence_enabled = False executor.kv_cache_transceiver.pipeline_transfer_enabled = True executor.model_engine.attn_runtime_features.chunked_prefill = True executor.dist.pp_size = 2 @@ -698,6 +700,7 @@ def test_pipelined_transfer_requires_single_beam_for_context_request(): from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor executor = MagicMock(_is_kv_manager_v2=True) + executor.model_engine._dspark_confidence_enabled = False executor.kv_cache_transceiver.pipeline_transfer_enabled = True executor.model_engine.attn_runtime_features.chunked_prefill = True @@ -718,6 +721,7 @@ def test_pipelined_transfer_allows_non_disaggregated_request(): from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor executor = MagicMock(_is_kv_manager_v2=True) + executor.model_engine._dspark_confidence_enabled = False executor.is_warmup = False executor.max_beam_width = 1 executor.kv_cache_transceiver.pipeline_transfer_enabled = True @@ -741,6 +745,7 @@ def test_pipelined_transfer_allows_generation_only_request(): from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor executor = MagicMock(_is_kv_manager_v2=True) + executor.model_engine._dspark_confidence_enabled = False executor.is_warmup = False executor.max_beam_width = 2 executor.kv_cache_transceiver.pipeline_transfer_enabled = True diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index a285cb43c864..c1587446ced3 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -5705,3 +5705,173 @@ 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 cost table.""" + kw["enable_confidence_scheduling"] = True + kw.setdefault("confidence_cost_table_path", "/tmp/costs.json") + return self._cfg(**kw) + + def test_defaults_are_off(self): + c = self._cfg() + assert c.enable_confidence_scheduling is False + assert c.confidence_min_predicted_gain == pytest.approx(0.04) + + def test_confidence_paths_are_ignored_when_scheduling_is_off(self): + # A static config copied from a dynamic one keeps working. + c = self._cfg(confidence_cost_table_path="/tmp/costs.json", + confidence_calibration_path="/tmp/calibration.json") + assert c.enable_confidence_scheduling is False + + def test_scheduling_requires_a_cost_table(self): + with pytest.raises(ValueError, match="confidence_cost_table_path"): + self._cfg(enable_confidence_scheduling=True) + + def test_min_predicted_gain_must_be_non_negative(self): + with pytest.raises(ValidationError, + match="confidence_min_predicted_gain"): + self._scheduled(confidence_min_predicted_gain=-0.01) + + @pytest.mark.parametrize("field", [ + "validate_confidence_artifact_freshness", + "confidence_sps_live_fingerprint_path", + "confidence_admission_receipt_path", + "confidence_admission_receipt_sha256", + "confidence_verify_len_tiers", + ]) + def test_removed_fields_are_rejected(self, field): + with pytest.raises(ValidationError, match=field): + self._scheduled(**{field: None}) + + def test_physical_k1_resolves_confidence_scheduling_to_static(self): + c = self._scheduled( + max_draft_len=1, + confidence_calibration_path="/tmp/calibration.json", + ) + + assert c.enable_confidence_scheduling is False + + 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 + + @staticmethod + def _environment(**overrides): + values = dict( + cuda_graph_config=SimpleNamespace(enable_padding=True), + moe_config=SimpleNamespace(load_balancer=None), + 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, + context_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"), + ({ + "context_parallel_size": 2 + }, "context_parallel_size=1"), + ( + { + "moe_config": + SimpleNamespace(load_balancer=SimpleNamespace( + mode="per_layer_dynamic")) + }, + "per-layer dynamic EPLB", + ), + ( + { + "sparse_attention_config": + SimpleNamespace( + enable_heuristic_topk=True, + use_cute_dsl_topk=False, + use_cute_dsl_paged_mqa_logits=False, + ) + }, + "enable_heuristic_topk", + ), + ( + { + "sparse_attention_config": + SimpleNamespace( + enable_heuristic_topk=False, + use_cute_dsl_topk=True, + use_cute_dsl_paged_mqa_logits=False, + ) + }, + "use_cute_dsl_topk", + ), + ( + { + "speculative_config": + SimpleNamespace(draft_is_embedded_in_target=True, + acceptance_rate_window_size=16, + acceptance_rate_threshold=0.3) + }, + "acceptance-rate gate", + ), + ], + ) + 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) + + @pytest.mark.parametrize("mode", [None, "standard"]) + def test_confidence_environment_accepts_padded_cuda_graphs(self, mode): + load_balancer = None if mode is None else SimpleNamespace(mode=mode) + TorchLlmArgs._validate_dspark_confidence_environment( + self._environment(moe_config=SimpleNamespace( + load_balancer=load_balancer)))