From 20f3c9c82f7ed5080abf9d4287308f350c13c785 Mon Sep 17 00:00:00 2001 From: Zhaoyuan He Date: Fri, 2 Oct 2026 12:44:46 -0700 Subject: [PATCH 1/2] [TRTLLM-14620][feat] Add DSpark confidence verification policy Validate exact measured-cost geometry and workload admission, allocate confidence-ranked verification prefixes, and keep host confidence ownership separate from completed sampler windows. Include focused policy, cost-table, allocation and admission behavior tests without enabling serving integration. Signed-off-by: Zhaoyuan He --- .../_torch/speculative/dspark_planner.py | 1226 +++++++++++++++++ .../_torch/speculative/dspark_schedule.py | 173 +++ .../_torch/speculative/dspark_verify.py | 395 ++++++ .../hw_agnostic/test_dspark_cost_table.py | 751 ++++++++++ .../hw_agnostic/test_dspark_policy.py | 294 ++++ .../hw_agnostic/test_dspark_schedule.py | 105 ++ .../test_dspark_workload_admission.py | 322 +++++ 7 files changed, 3266 insertions(+) create mode 100644 tensorrt_llm/_torch/speculative/dspark_planner.py create mode 100644 tensorrt_llm/_torch/speculative/dspark_schedule.py create mode 100644 tensorrt_llm/_torch/speculative/dspark_verify.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_table.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_policy.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_schedule.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_workload_admission.py diff --git a/tensorrt_llm/_torch/speculative/dspark_planner.py b/tensorrt_llm/_torch/speculative/dspark_planner.py new file mode 100644 index 000000000000..931dc1e94554 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_planner.py @@ -0,0 +1,1226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Authenticated DSpark verification costs and exact ``(G, V)`` selection.""" + +import hashlib +import json +import math +from collections.abc import Sequence +from dataclasses import dataclass, field +from pathlib import Path + +import numpy as np + +_MAX_EXACT_COMPACT_CELLS_PER_G = 4 +_MAX_EXACT_COMPACT_CELLS_TOTAL = 32 + +__all__ = [ + "ConfidenceWorkloadAdmission", + "EXACT_SPS_SELECTOR_IDENTITY_SHA256", + "ExactSpsDrainGuard", + "ExactSpsCostRow", + "ExactSpsCostTable", + "evaluate_confidence_workload_admission", + "load_confidence_workload_admission", + "load_runtime_sps_cost_table", + "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``.""" + + token_counts: Sequence[int] + step_time_ms: Sequence[float] + + def __post_init__(self) -> None: + if len(self.token_counts) != len(self.step_time_ms): + raise ValueError( + f"token_counts ({len(self.token_counts)}) and step_time_ms " + f"({len(self.step_time_ms)}) must have the same length" + ) + if not self.token_counts: + raise ValueError("ExactSpsCostRow requires at least one measured point") + token_counts = tuple( + _require_exact_int(value, field="measured verifier budget", minimum=0) + 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 + ) + if any(b <= a for a, b in zip(token_counts, token_counts[1:])): + raise ValueError("token_counts must be strictly increasing") + object.__setattr__(self, "token_counts", token_counts) + 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)``. + + ``G`` is the rank-local padded CUDA-graph batch size and ``V`` is the + rank-local submitted ragged verifier-token count, including anchor and + pad-row tokens. ``V=0`` is a sentinel for the native static physical-K path: it is + the mandatory fallback comparator and is never a ragged capture bucket. + Every positive cell must fit the physical verifier window + ``G <= V <= G * (max_draft_len + 1)``. + No interpolation is permitted on either axis: an unmeasured graph shape + is a configuration error rather than an estimate. + """ + + tables: dict[int, ExactSpsCostRow] + max_draft_len: int + minimum_predicted_gain: float = 0.01 + iteration_drain_guard: ExactSpsDrainGuard | None = None + identity_sha256: str = field(init=False) + + def __post_init__(self) -> None: + max_draft_len = _require_exact_int(self.max_draft_len, field="max_draft_len", minimum=1) + normalized: dict[int, ExactSpsCostRow] = {} + for graph_batch_size, table in self.tables.items(): + canonical_graph_batch_size = _require_exact_int( + graph_batch_size, field="graph batch size", minimum=1 + ) + if canonical_graph_batch_size in normalized: + raise ValueError("ExactSpsCostTable contains duplicate canonical graph batch sizes") + normalized[canonical_graph_batch_size] = table + if not normalized: + raise ValueError("ExactSpsCostTable requires at least one graph batch size") + 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 + ) + 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}" + ) + max_verifier_budget = graph_batch_size * (max_draft_len + 1) + invalid_positive_budgets = [ + verifier_budget + for verifier_budget in budgets[1:] + if (verifier_budget < graph_batch_size or verifier_budget > max_verifier_budget) + ] + if invalid_positive_budgets: + raise ValueError( + "Exact SPS 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}" + ) + production_budgets = [ + verifier_budget + for verifier_budget in budgets[1:] + if verifier_budget < max_verifier_budget + ] + if len(production_budgets) > _MAX_EXACT_COMPACT_CELLS_PER_G: + raise ValueError( + "Exact SPS graph budget exceeds the production limit of " + f"{_MAX_EXACT_COMPACT_CELLS_PER_G} compact V cells per G; " + f"G={graph_batch_size} has {len(production_budgets)}" + ) + production_cell_count = sum( + len( + [ + verifier_budget + for verifier_budget in table.token_counts[1:] + if verifier_budget < graph_batch_size * (max_draft_len + 1) + ] + ) + for graph_batch_size, table in normalized.items() + ) + 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}" + ) + object.__setattr__(self, "tables", normalized) + object.__setattr__(self, "max_draft_len", max_draft_len) + minimum_predicted_gain = _require_nonnegative_finite_number( + 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 + # number of cells must not be allowed to index each other's yields. + identity_payload = { + "schema": "exact-sps-runtime-v1", + "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, + "token_counts": [int(value) for value in table.token_counts], + "step_time_ms": [float(value) for value in table.step_time_ms], + } + for graph_batch_size, table in sorted(normalized.items()) + ], + } + object.__setattr__(self, "identity_sha256", _canonical_json_sha256(identity_payload)) + + @property + def collective_identity_words(self) -> tuple[int, ...]: + """Full SHA256 as eight transport-safe unsigned 32-bit integers.""" + return tuple( + int(self.identity_sha256[offset : offset + 8], 16) + for offset in range(0, len(self.identity_sha256), 8) + ) + + def for_graph_batch_size(self, num_requests: int) -> ExactSpsCostRow: + """Return the directly measured V cells for one exact G.""" + graph_batch_size = _require_exact_int(num_requests, field="graph batch size", minimum=1) + try: + 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}" + ) from error + + def step_times(self, num_tokens: np.ndarray, num_requests: int) -> np.ndarray: + """Return direct whole-step measurements without interpolation.""" + table = self.for_graph_batch_size(num_requests) + measured = { + _require_exact_int(token_count, field="measured verifier budget", minimum=0): float( + step_time + ) + for token_count, step_time in zip(table.token_counts, table.step_time_ms) + } + requested = np.asarray(num_tokens) + requested_budgets = [ + _require_exact_int(value, field="requested verifier budget", minimum=0) + for value in requested.reshape(-1) + ] + 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}" + ) + return np.asarray( + [measured[value] for value in requested_budgets], + dtype=np.float64, + ).reshape(requested.shape) + + def step_time(self, num_tokens: int, num_requests: int) -> float: + """Return one directly measured ``T(G, V)`` value.""" + return float(self.step_times(np.asarray([num_tokens]), num_requests)[0]) + + 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 + ) + if include_native: + return budgets + return tuple(value for value in budgets if value != 0) + + def candidate_cells(self) -> tuple[tuple[int, int], ...]: + """Stable positive ``(G,V)`` ordering for collective payloads.""" + return tuple( + (graph_batch_size, verifier_budget) + for graph_batch_size in sorted(self.tables) + for verifier_budget in self.production_candidate_budgets(graph_batch_size) + ) + + def production_candidate_budgets(self, num_requests: int) -> tuple[int, ...]: + """Compact serving cells, excluding native and full-token controls.""" + graph_batch_size = _require_exact_int(num_requests, field="graph batch size", minimum=1) + full_budget = graph_batch_size * (self.max_draft_len + 1) + return tuple( + verifier_budget + for verifier_budget in self.candidate_budgets(graph_batch_size) + if verifier_budget < full_budget + ) + + +def _read_sps_cost_payload(path: str | Path) -> dict[str, object]: + with Path(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") + 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.""" + payload = _read_sps_cost_payload(path) + schema_version = payload.get("schema_version") + if type(schema_version) is not int or schema_version != 2: + raise ValueError("Exact SPS cost artifacts require schema_version=2") + if live_engine_fingerprint_path is None: + raise ValueError( + "Schema-v2 exact SPS artifacts require an independently generated " + "live engine fingerprint path" + ) + live_engine_fingerprint = _read_sps_cost_payload(live_engine_fingerprint_path) + validate_sps_cost_table_payload( + payload, + graph_batch_sizes=graph_batch_sizes, + max_draft_len=max_draft_len, + live_engine_fingerprint=live_engine_fingerprint, + ) + return _build_exact_sps_cost_table(payload, max_draft_len=max_draft_len), payload + + +def _build_exact_sps_cost_table( + payload: dict[str, object], + *, + max_draft_len: int, +) -> 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", + ), + ) + return ExactSpsCostTable( + tables=parsed_tables, + 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, + ) + + +def _canonical_json_sha256(value: object) -> str: + encoded = json.dumps(value, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(encoded).hexdigest() + + +_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, +) -> 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 + + +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 _parse_graph_batch_size_key(value: object) -> int: + if ( + not isinstance(value, str) + or not value.isascii() + or not value.isdecimal() + or value.startswith("0") + ): + raise TypeError( + f"SPS 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 + + +def _require_json_int(value: object, *, field: str, minimum: int) -> int: + if type(value) is not int: + raise TypeError(f"{field} must be a JSON integer, got {value!r}") + if value < minimum: + raise ValueError(f"{field} must be >= {minimum}, got {value}") + return value + + +def _require_exact_int(value: object, *, field: str, minimum: int) -> int: + if isinstance(value, bool) or not isinstance(value, (int, np.integer)): + raise TypeError(f"{field} must be an integer, got {value!r}") + parsed = int(value) + if parsed < minimum: + raise ValueError(f"{field} must be >= {minimum}, got {parsed}") + return parsed + + +def _require_positive_finite_number(value: object, *, field: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError(f"{field} must be a JSON number, got {value!r}") + parsed = float(value) + if not math.isfinite(parsed) or parsed <= 0.0: + raise ValueError(f"{field} must be positive and finite, got {value!r}") + return parsed + + +def _require_nonnegative_finite_number(value: object, *, field: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError(f"{field} must be a JSON number, got {value!r}") + parsed = float(value) + if not math.isfinite(parsed) or parsed < 0.0: + raise ValueError(f"{field} must be non-negative and finite, got {value!r}") + 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 _sha256_file(path: str | Path) -> str: + digest = hashlib.sha256() + with Path(path).open("rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _read_strict_json_object(path: str | Path, *, 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 + + with Path(path).open(encoding="utf-8") as file: + payload = json.load(file, 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_sha256 = _sha256_file(path) + if receipt_sha256 != expected_receipt_sha256: + raise ValueError("Confidence workload admission receipt does not match its pinned SHA256") + envelope = _validate_exact_fields( + _read_strict_json_object(path, 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") + if admission["sps_cost_table_sha256"] != _sha256_file(sps_cost_table_path): + raise ValueError("Confidence workload admission SPS cost-table SHA256 does not match") + if admission["live_engine_fingerprint_sha256"] != _sha256_file(live_engine_fingerprint_path): + raise ValueError( + "Confidence workload admission live-engine fingerprint SHA256 does not match" + ) + + live_fingerprint = _validate_exact_fields( + _read_strict_json_object(live_engine_fingerprint_path, 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 = _read_strict_json_object(sps_cost_table_path, 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, + ) + 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 + """ + 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) + native_budgets = cost_table.candidate_budgets(graph_batch_size, include_native=True) + if not native_budgets or native_budgets[0] != 0: + raise ValueError(f"Exact SPS G={graph_batch_size} has no V=0 native fallback") + if any(type(value) is not int for value in compact_expected_yields): + raise TypeError("compact_expected_yields keys must be integer V values") + if 0 in compact_expected_yields: + raise ValueError("compact_expected_yields must not include native V=0") + if set(compact_expected_yields) != set(measured_budgets): + 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" + ) + if any(value > native_yield for value in compact_yields.values()): + raise ValueError("compact expected yield must not exceed native expected yield") + native_score = native_yield / cost_table.step_time(0, graph_batch_size) + compact_scores = { + 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 + for budget in measured_budgets + if compact_scores[budget] > native_score and compact_scores[budget] >= required_score + ] + 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, + ), + ) + ) diff --git a/tensorrt_llm/_torch/speculative/dspark_schedule.py b/tensorrt_llm/_torch/speculative/dspark_schedule.py new file mode 100644 index 000000000000..3a12d6d27bf1 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_schedule.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""DSpark confidence-scheduled verification (arXiv:2607.05147 ยง5). + +``survival[r][j] = prod(conf[r][:j+1])`` ranks every ``(request, position)`` +candidate globally; the best ``budget`` are admitted. Survival is +non-increasing along each request's positions, so any global top-k is +automatically a per-request prefix and the allocation only counts admitted +candidates per row. Only the number of verified tokens is decided here, never +acceptance, so scheduling is lossless w.r.t. the target distribution. Ties are +broken by ``(position, request)``. Ranks with different confidence values can +produce different allocations; cross-rank agreement belongs to the caller. +""" + +from dataclasses import dataclass + +import torch + +__all__ = [ + "DSparkScheduleConfig", + "compute_survival", + "schedule_verify_lens_topk", +] + + +# Raw logit for confidence rows that carry no measurement: sigmoid(30/T) ~ 1.0 +# at any fitted temperature, so an unknown row is treated as "certainly accept" +# -- optimistic by design. Lives here because the worker (writes it) and the +# planner (counts it) cannot import each other. +NEUTRAL_CONFIDENCE_LOGIT = 30.0 + + +@dataclass(frozen=True) +class DSparkScheduleConfig: + """Bounds for the per-request verify length. + + Attributes: + block_size: draft block length ``K``. + min_verify_len: floor on every request's draft verify length; must + stay >= 1. Budget is allocated *above* this floor. The target + anchor contributes a separate token to the verifier input. + max_verify_len: cap on every request's verify length; 0 means + ``block_size``; never exceeds ``block_size``. + survival_eps: numerical floor (not a tuning knob); candidates below it + are dropped before ranking. + """ + + block_size: int + min_verify_len: int = 1 + max_verify_len: int = 0 + survival_eps: float = 1e-6 + + def __post_init__(self) -> None: + if self.block_size < 1: + raise ValueError(f"block_size must be >= 1, got {self.block_size}") + if self.min_verify_len < 1: + raise ValueError( + f"min_verify_len must be >= 1 (every request verifies its anchor " + f"token), got {self.min_verify_len}" + ) + if self.resolved_max_verify_len < self.min_verify_len: + raise ValueError( + f"max_verify_len ({self.resolved_max_verify_len}) < min_verify_len " + f"({self.min_verify_len})" + ) + if not 0.0 <= self.survival_eps < 1.0: + raise ValueError(f"survival_eps must be in [0, 1), got {self.survival_eps}") + + @property + def resolved_max_verify_len(self) -> int: + """``max_verify_len`` with 0 meaning "the whole block", clamped to ``block_size``.""" + cap = self.max_verify_len or self.block_size + return min(int(cap), int(self.block_size)) + + @property + def schedulable_per_request(self) -> int: + """Positions a single request can win from the budget, above the floor.""" + return self.resolved_max_verify_len - self.min_verify_len + + +def compute_survival(confidence: torch.Tensor) -> torch.Tensor: + """``[bs, K]`` conditional acceptance probabilities -> prefix survival. + + ``out[r][j] = prod_{i <= j} confidence[r][i]``, computed in fp32. The result + is non-increasing along ``j``, which :func:`schedule_verify_lens_topk` relies + on to keep its allocation prefix-shaped. + """ + if confidence.dim() != 2: + raise ValueError(f"confidence must be [bs, K], got shape {tuple(confidence.shape)}") + return torch.cumprod(confidence.to(torch.float32), dim=1) + + +def schedule_verify_lens_topk( + *, + survival: torch.Tensor, + budget: int | torch.Tensor, + cfg: DSparkScheduleConfig, +) -> torch.Tensor: + """Allocate ``budget`` verify tokens across the batch by global survival rank. + + Args: + survival: ``[bs, K]`` prefix-survival probabilities from + :func:`compute_survival`. + budget: total verify tokens to hand out *above* the per-request floor. + A 0-d integer tensor is honoured without a host sync (the value + stays device-resident through the ``rank < budget`` cut), which is + what lets the whole function run inside a captured graph. + cfg: bounds; see :class:`DSparkScheduleConfig`. + Returns: + ``[bs]`` int32 verify lengths, each in + ``[min_verify_len, resolved_max_verify_len]``, summing to at most + ``bs * min_verify_len + budget``. + + Ties use position before request. This determines the allocation for equal + survival values, but does not remove differences between rank-local inputs. + """ + if survival.dim() != 2: + raise ValueError(f"survival must be [bs, K], got shape {tuple(survival.shape)}") + bs, num_positions = survival.shape + if num_positions != cfg.block_size: + raise ValueError( + f"survival has {num_positions} positions but block_size is {cfg.block_size}" + ) + + device = survival.device + floor = int(cfg.min_verify_len) + schedulable = int(cfg.schedulable_per_request) + verify_lens = torch.full((bs,), floor, dtype=torch.int32, device=device) + if not isinstance(budget, torch.Tensor): + budget = int(budget) + if budget <= 0: + return verify_lens + if bs == 0 or schedulable <= 0: + return verify_lens + + # Positions [0, floor) are already granted to every request; only those + # above compete for budget. + candidates = survival[:, floor : floor + schedulable].to(torch.float32) + flat = candidates.reshape(-1) + eligible = flat >= cfg.survival_eps + + # Sort the tie-break key first, then stable-sort on the scores: the final + # order is a pure function of (survival, position, request). Ineligible + # candidates score -1 and are masked by ``& eligible`` below; no + # ``.item()``/host sync anywhere, so the path stays capture-safe. + scores = torch.where(eligible, flat, torch.full_like(flat, -1.0)) + tie_break = torch.arange(schedulable, device=device).repeat(bs) * bs + torch.arange( + bs, device=device + ).repeat_interleave(schedulable) + order = torch.argsort(tie_break, stable=True) + order = order[torch.argsort(scores[order], descending=True, stable=True)] + + rank = torch.empty_like(order) + rank[order] = torch.arange(order.numel(), device=device) + selected = (rank < budget) & eligible + + # survival is non-increasing along positions, so the chosen set for a request + # is necessarily a prefix -- counting per row recovers the full allocation. + granted = selected.view(bs, schedulable).sum(dim=1).to(torch.int32) + verify_lens += granted + return torch.clamp(verify_lens, min=floor, max=cfg.resolved_max_verify_len) diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py new file mode 100644 index 000000000000..69d5b707592d --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -0,0 +1,395 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Rank-local confidence staging and exact DSpark ``(G, V)`` policy inputs.""" + +from collections.abc import Callable, Sequence +from dataclasses import dataclass, replace + +import numpy as np +import torch + +from ..._utils import prefer_pinned +from .dspark_planner import ExactSpsCostTable +from .dspark_schedule import DSparkScheduleConfig, compute_survival, schedule_verify_lens_topk + + +@dataclass(frozen=True) +class ExactSpsLocalDecision: + """One rank's pre-collective yield curve for exact ``(G,V)`` cells.""" + + num_requests: int + survival: torch.Tensor + native_expected_yield: float + compact_expected_yields: tuple[float, ...] + + +def _exact_cell_geometry( + *, + num_real: int, + graph_batch_size: int, + verifier_budget: int, + min_verify_len: int, + max_verify_len: int, +) -> tuple[int, int, int] | None: + """Return ``(pad_tokens, real_tokens, extra_budget)`` or infeasible. + + Real-containing ranks use one shared pad-row length. A zero-real rank may + use the deterministic quotient/remainder split ``V = G*q + r`` backed by + two cached dummy objects (low=q, high=q+1); the returned pad length is q. + The same calculation is reused by yield modeling and execution so the + policy cannot price one layout and run another. + """ + n = int(num_real) + graph_batch_size = int(graph_batch_size) + verifier_budget = int(verifier_budget) + if n < 0 or n > graph_batch_size or graph_batch_size <= 0: + return None + if n == 0: + # Row zero is the scheduled ADP dummy; CUDA padding contributes the + # remaining G-1 rows. Two cached CUDA dummy variants represent the + # quotient/remainder split without allocating G distinct KV objects. + pad_tokens, remainder = divmod(verifier_budget, graph_batch_size) + if pad_tokens < 1 or pad_tokens + int(remainder > 0) > int(max_verify_len) + 1: + return None + return pad_tokens, 0, 0 + num_pad = graph_batch_size - n + real_floor = n * (int(min_verify_len) + 1) + real_capacity = n * (int(max_verify_len) + 1) + if num_pad == 0: + if not real_floor <= verifier_budget <= real_capacity: + return None + real_tokens = verifier_budget + return 0, real_tokens, real_tokens - real_floor + + pad_capacity = int(max_verify_len) + 1 + pad_lo = max(1, -(-(verifier_budget - real_capacity) // num_pad)) + pad_hi = min(pad_capacity, (verifier_budget - real_floor) // num_pad) + if pad_lo > pad_hi: + return None + pad_tokens = pad_lo + real_tokens = verifier_budget - num_pad * pad_tokens + if not real_floor <= real_tokens <= real_capacity: + return None + return pad_tokens, real_tokens, real_tokens - real_floor + + +__all__ = ["DSparkVerifyPlanner"] + + +class DSparkVerifyPlanner: + """Stages confidence and prepares exact ``(G,V)`` policy inputs. + + Args: + cfg: verify-length bounds. + cost_table: authenticated exact ``T(G,V)`` table installed by + ``ModelEngine`` before the first decision. + apply_calibration: maps raw confidence logits to probabilities. Normally + ``confidence_head.apply_sts``; defaults to a plain sigmoid. + device_windows: use the lag-2 snapshot for capacity while a device + prologue ranks the current requests. + """ + + def __init__( + self, + *, + cfg: DSparkScheduleConfig, + cost_table: ExactSpsCostTable | None = None, + apply_calibration: Callable[[torch.Tensor], torch.Tensor] | None = None, + device_windows: bool = False, + ) -> None: + self.cfg = cfg + self.cost_table = cost_table + self.apply_calibration = apply_calibration or torch.sigmoid + + self._host_buffer: torch.Tensor | None = None + self._copy_event: torch.cuda.Event | None = None + self._snapshot_valid = False + + # Device-window mode uses the older confidence copy for capacity and + # ranked host fallback. A downstream device prologue can rerank from + # current confidence. Each copy is read only after its event completes. + self.device_windows = bool(device_windows) + self._prev_buffer: torch.Tensor | None = None + self._prev_event: torch.cuda.Event | None = None + self._prev_valid = False + # Count every path that declines to trim; degradation here is silent. + self.stats = { + "fallback_no_snapshot": 0, + "fallback_no_confidence": 0, + "fallback_short_snapshot": 0, + } + + def install_runtime_cost_table(self, cost_table: ExactSpsCostTable) -> None: + """Install ModelEngine's already validated runtime cost object.""" + if not isinstance(cost_table, ExactSpsCostTable): + raise TypeError("DSpark confidence scheduling requires an exact SPS cost table") + if cost_table.max_draft_len != self.cfg.resolved_max_verify_len: + raise ValueError( + "Exact SPS max_draft_len does not match planner: " + f"table={cost_table.max_draft_len}, " + f"planner={self.cfg.resolved_max_verify_len}" + ) + if self.cost_table is not None and self.cost_table != cost_table: + raise RuntimeError( + "DSpark planner already holds a different SPS cost object; " + "graph capture and runtime scheduling must share ModelEngine's " + "single validated instance" + ) + self.cost_table = cost_table + + @property + def exact_cost_table(self) -> ExactSpsCostTable | None: + return self.cost_table + + def prepare_exact_sps_decision( + self, + *, + num_gen_requests: int, + rows: Sequence[int] | None = None, + ) -> ExactSpsLocalDecision | None: + """Price every measured V locally before the existing allgather. + + The payload includes the whole measured grid because the common G is + not known until attention-DP ranks exchange their row counts. Pad rows + have no expected output and initially consume one anchor token; any V + above real-row capacity is feasible but adds no yield. + """ + table = self.exact_cost_table + if table is None: + raise RuntimeError("prepare_exact_sps_decision requires an exact SPS table") + n = int(num_gen_requests) + if n < 0: + return None + + if n == 0: + # An idle ADP rank has no confidence rows and contributes no + # expected output. It still advertises every pad-only cell that the + # low/high dummy pair can represent; resource readiness is agreed + # by the executor before a selected cell can execute. + max_verify_len = int(self.cfg.resolved_max_verify_len) + compact_expected_yields = tuple( + 0.0 + if _exact_cell_geometry( + num_real=0, + graph_batch_size=graph_batch_size, + verifier_budget=verifier_budget, + min_verify_len=self.cfg.min_verify_len, + max_verify_len=max_verify_len, + ) + is not None + else -1.0 + for graph_batch_size, verifier_budget in table.candidate_cells() + ) + return ExactSpsLocalDecision( + num_requests=0, + survival=torch.empty((0, max_verify_len), dtype=torch.float32), + native_expected_yield=0.0, + compact_expected_yields=compact_expected_yields, + ) + + if self.device_windows: + snapshot = self._ready_older_snapshot() + if snapshot is None: + return None + selected = self._gather_rows( + num_gen_requests=n, + rows=rows, + snapshot=snapshot, + ) + else: + selected = self._gather_rows( + num_gen_requests=n, + rows=rows, + ) + if selected is None: + return None + survival = compute_survival(self.apply_calibration(selected)) + floor = int(self.cfg.min_verify_len) + max_verify_len = int(self.cfg.resolved_max_verify_len) + # Build every exact-cell yield from one float64 decomposition. Mixing + # torch float32 native/base reductions with a NumPy float64 optional + # prefix can make a saturated compact cell appear better than native + # solely because of reduction order. + survival_values = survival.numpy().astype(np.float64, copy=False) + base_expected_yield = float(n) + float(survival_values[:, :floor].sum(dtype=np.float64)) + candidates = np.sort(survival_values[:, floor:max_verify_len].reshape(-1))[::-1] + prefix = np.concatenate(([0.0], np.cumsum(candidates, dtype=np.float64))) + native_expected_yield = ( + base_expected_yield + + float(prefix[-1]) + + float(survival_values[:, max_verify_len:].sum(dtype=np.float64)) + ) + compact_expected_yields = [] + for graph_batch_size, verifier_budget in table.candidate_cells(): + geometry = _exact_cell_geometry( + num_real=n, + graph_batch_size=graph_batch_size, + verifier_budget=verifier_budget, + min_verify_len=floor, + max_verify_len=max_verify_len, + ) + if geometry is None: + compact_expected_yields.append(-1.0) + continue + _, _, extra = geometry + compact_expected_yields.append(base_expected_yield + float(prefix[extra])) + return ExactSpsLocalDecision( + num_requests=n, + survival=survival, + native_expected_yield=native_expected_yield, + compact_expected_yields=tuple(compact_expected_yields), + ) + + def allocate_exact_sps_candidate( + self, + decision: ExactSpsLocalDecision, + *, + graph_batch_size: int, + verifier_budget: int, + ) -> tuple[list[int], int, int]: + """Allocate one globally selected exact V on this rank. + + Returns per-request draft windows, the real-row budget above the + mandatory floor, and the shared pad-row token length. The bucket fitter + receives all three parts of the exact plan and may only validate them. + """ + table = self.exact_cost_table + if table is None: + raise RuntimeError("allocate_exact_sps_candidate requires an exact SPS table") + n = int(decision.num_requests) + graph_batch_size = int(graph_batch_size) + verifier_budget = int(verifier_budget) + if verifier_budget not in table.production_candidate_budgets(graph_batch_size): + raise ValueError(f"Unmeasured exact SPS cell G={graph_batch_size}, V={verifier_budget}") + floor = int(self.cfg.min_verify_len) + max_verify_len = int(self.cfg.resolved_max_verify_len) + geometry = _exact_cell_geometry( + num_real=n, + graph_batch_size=graph_batch_size, + verifier_budget=verifier_budget, + min_verify_len=floor, + max_verify_len=max_verify_len, + ) + if geometry is None: + raise ValueError( + f"Exact SPS cell G={graph_batch_size}, V={verifier_budget} " + f"cannot fit {n} real requests with verify floor {floor}" + ) + pad_tokens, real_tokens, budget = geometry + if n == 0: + if real_tokens != 0 or budget != 0: + raise RuntimeError("zero-real exact SPS geometry assigned real-row work") + return [], 0, int(pad_tokens) + exact_cfg = replace(self.cfg, survival_eps=0.0) + lens = schedule_verify_lens_topk( + survival=decision.survival, budget=budget, cfg=exact_cfg + ).tolist() + if sum(int(value) + 1 for value in lens) != real_tokens: + raise RuntimeError("Exact SPS allocation did not spend its real-row token target") + return [int(value) for value in lens], int(budget), int(pad_tokens) + + @property + def max_verify_len(self) -> int: + return self.cfg.resolved_max_verify_len + + def stage_confidence(self, confidence_logits: torch.Tensor) -> None: + """Start a non-blocking device->host copy of the confidence buffer. + + ``confidence_logits`` is the worker's whole slot-indexed buffer, staged + in full so exact-policy preparation can look up any live request by + slot across the one-iteration lag. Never synchronizes: the copy is + issued on the current stream and an event is recorded; policy + preparation polls the event and skips a snapshot that has not landed. + The caller must order confidence writes and staging on that stream. + These buffers hold logits, not completed-step verification lengths. + """ + if confidence_logits is None or confidence_logits.shape[0] == 0: + return + if self.device_windows: + # Rotate before overwriting: what was current becomes the older + # (lag-2) snapshot the budget argmax reads; the old older buffer's + # storage is reused for this staging. Pure reference swaps. + self._host_buffer, self._prev_buffer = (self._prev_buffer, self._host_buffer) + self._copy_event, self._prev_event = (self._prev_event, self._copy_event) + self._prev_valid = self._snapshot_valid + self._snapshot_valid = False + if self._host_buffer is None or self._host_buffer.shape != confidence_logits.shape: + self._host_buffer = torch.empty( + confidence_logits.shape, + dtype=torch.float32, + device="cpu", + pin_memory=prefer_pinned(), + ) + self._copy_event = torch.cuda.Event() + self._host_buffer.copy_(confidence_logits, non_blocking=True) + self._copy_event.record() + self._snapshot_valid = True + + def _ready_snapshot(self) -> torch.Tensor | None: + if not self._snapshot_valid or self._host_buffer is None: + return None + # query() is non-blocking; a not-yet-landed copy just means this step + # decides without it. + if self._copy_event is not None and not self._copy_event.query(): + return None + return self._host_buffer + + def _ready_older_snapshot(self) -> torch.Tensor | None: + """The lag-2 snapshot (device-window mode's capacity input), or None. + + The event must have completed before the host reads the buffer. + """ + if not self._prev_valid or self._prev_buffer is None: + return None + if self._prev_event is not None and not self._prev_event.query(): + return None + return self._prev_buffer + + def _gather_rows( + self, + *, + num_gen_requests: int, + rows: Sequence[int] | None, + snapshot: torch.Tensor | None = None, + ) -> torch.Tensor | None: + """This step's confidence, one row per generation request, or None. + + ``rows`` is the per-request buffer row index (``DSparkWorker. + confidence_row_for``); the batch is reshuffled across the one-iteration + lag, so positional reads (``rows=None``) are only correct for callers + that own the ordering. Declines (None) rather than returning fewer + rows than the batch, which would partially assign verify windows + downstream. ``snapshot`` overrides the default (current) snapshot when + the caller already picked one (the lag-2 budget read). + """ + if snapshot is None: + snapshot = self._ready_snapshot() + if snapshot is None: + self.stats["fallback_no_snapshot"] += 1 + return None + if rows is None: + if snapshot.shape[0] < num_gen_requests: + self.stats["fallback_short_snapshot"] += 1 + return None + selected = snapshot[:num_gen_requests] + else: + if len(rows) != num_gen_requests: + self.stats["fallback_short_snapshot"] += 1 + return None + selected = snapshot[torch.as_tensor(list(rows), dtype=torch.long)] + if selected.numel() == 0: + self.stats["fallback_no_confidence"] += 1 + return None + return selected 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 new file mode 100644 index 000000000000..e4a35388b097 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_table.py @@ -0,0 +1,751 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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.dspark_planner import ( + ExactSpsCostRow, + ExactSpsCostTable, + ExactSpsDrainGuard, + load_runtime_sps_cost_table, + 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, + ) + + +def _cfg(**kw) -> DSparkScheduleConfig: + return DSparkScheduleConfig(block_size=BLOCK, **kw) + + +# -------------------------------------------------------------------------- +# survival +# -------------------------------------------------------------------------- + + +def test_survival_is_cumulative_product(): + conf = torch.tensor([[0.9, 0.8, 0.5, 1.0, 1.0, 1.0, 1.0]]) + surv = compute_survival(conf) + 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 = { + "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 _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_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) + + 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) + + +def test_exact_cost_table_rejects_fractional_programmatic_cells(): + table = ExactSpsCostRow(token_counts=(0, 704), step_time_ms=(8.0, 7.0)) + with pytest.raises(TypeError, match="graph batch size"): + ExactSpsCostTable(tables={128.5: table}, max_draft_len=5) + with pytest.raises(TypeError, match="measured verifier budget"): + ExactSpsCostTable( + tables={128: ExactSpsCostRow(token_counts=(0, 704.5), step_time_ms=(8.0, 7.0))}, + max_draft_len=5, + ) + with pytest.raises(ValueError, match="V=0 native static physical-K"): + ExactSpsCostTable( + tables={128: ExactSpsCostRow(token_counts=(704,), step_time_ms=(7.0,))}, + max_draft_len=5, + ) + + +def test_exact_cost_table_identity_covers_grid_costs_and_policy(): + tables = { + 64: ExactSpsCostRow(token_counts=(0, 352), step_time_ms=(5.0, 4.5)), + 128: ExactSpsCostRow(token_counts=(0, 704), step_time_ms=(8.0, 7.0)), + } + first = ExactSpsCostTable(tables=tables, max_draft_len=5, minimum_predicted_gain=0.02) + reordered = ExactSpsCostTable( + tables=dict(reversed(tuple(tables.items()))), + max_draft_len=5, + minimum_predicted_gain=0.02, + ) + changed_cost = ExactSpsCostTable( + tables={ + **tables, + 128: ExactSpsCostRow(token_counts=(0, 704), step_time_ms=(8.0, 7.1)), + }, + max_draft_len=5, + 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 + assert len(first.collective_identity_words) == 8 + 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]) +def test_exact_cost_table_direct_construction_bounds_positive_cells(invalid_budget): + with pytest.raises(ValueError, match=r"G <= V <= G\*\(K\+1\)"): + ExactSpsCostTable( + tables={ + 128: ExactSpsCostRow(token_counts=(0, invalid_budget), step_time_ms=(8.0, 7.0)) + }, + max_draft_len=5, + ) + + +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], + 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, + ) + + +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], + 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], + 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) + + +def test_exact_selector_keeps_native_on_tie_or_below_threshold(): + table = 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(), + ) + + 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, + ) + == 0 + ) + 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 + ) + with pytest.raises(ValueError, match="must not include native V=0"): + select_exact_sps_candidate( + 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): + table = ExactSpsCostTable( + tables={ + graph_batch_size: ExactSpsCostRow( + 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, + ), + ) + + # 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. + assert ( + select_exact_sps_candidate( + graph_batch_size=graph_batch_size, + native_expected_yield=10.0, + compact_expected_yields={ + 3 * graph_batch_size: 6.2, + 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 + ) + + +def test_exact_selector_fails_closed_without_measured_compact_cell(): + graph_batch_size = 32 + 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, + 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_policy.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_policy.py new file mode 100644 index 000000000000..a5f7ddbc35f1 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_policy.py @@ -0,0 +1,294 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hardware-independent tests for rank-local DSpark confidence planning.""" + +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark_planner import ExactSpsCostRow, ExactSpsCostTable +from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig +from tensorrt_llm._torch.speculative.dspark_verify import ( + DSparkVerifyPlanner, + ExactSpsLocalDecision, + _exact_cell_geometry, +) + + +def _config() -> DSparkScheduleConfig: + return DSparkScheduleConfig(block_size=5, min_verify_len=1) + + +def _publish_snapshot(planner: DSparkVerifyPlanner, logits: torch.Tensor) -> None: + planner._host_buffer = logits + planner._copy_event = None + planner._snapshot_valid = True + + +def test_runtime_cost_table_is_authoritative() -> None: + planner = DSparkVerifyPlanner(cfg=_config()) + authoritative = ExactSpsCostTable( + tables={4: ExactSpsCostRow(token_counts=(0, 14), step_time_ms=(8.0, 7.0))}, + max_draft_len=5, + ) + + planner.install_runtime_cost_table(authoritative) + planner.install_runtime_cost_table(authoritative) + assert planner.cost_table is authoritative + + different = ExactSpsCostTable( + tables={4: ExactSpsCostRow(token_counts=(0, 14), step_time_ms=(8.0, 7.1))}, + max_draft_len=5, + ) + with pytest.raises(RuntimeError, match="different SPS cost object"): + planner.install_runtime_cost_table(different) + + +def test_runtime_cost_table_must_match_the_draft_block() -> None: + planner = DSparkVerifyPlanner(cfg=_config()) + incompatible = ExactSpsCostTable( + tables={4: ExactSpsCostRow(token_counts=(0, 12), step_time_ms=(8.0, 7.0))}, + max_draft_len=4, + ) + + with pytest.raises(ValueError, match="max_draft_len does not match"): + planner.install_runtime_cost_table(incompatible) + + +@pytest.mark.parametrize( + ("num_real", "graph_batch_size", "verifier_budget", "expected"), + [ + (100, 128, 704, (4, 592, 392)), + (0, 128, 704, (5, 0, 0)), + (4, 4, 24, (0, 24, 16)), + (4, 4, 25, None), + (5, 4, 14, None), + ], +) +def test_exact_cell_geometry_matches_the_executed_layout( + num_real: int, + graph_batch_size: int, + verifier_budget: int, + expected: tuple[int, int, int] | None, +) -> None: + assert ( + _exact_cell_geometry( + num_real=num_real, + graph_batch_size=graph_batch_size, + verifier_budget=verifier_budget, + min_verify_len=1, + max_verify_len=5, + ) + == expected + ) + + +def test_idle_attention_dp_rank_advertises_only_feasible_cells() -> None: + table = ExactSpsCostTable( + tables={ + 4: ExactSpsCostRow( + token_counts=(0, 12, 14, 24), + step_time_ms=(8.0, 6.0, 6.5, 8.1), + ) + }, + max_draft_len=5, + ) + planner = DSparkVerifyPlanner(cfg=_config(), cost_table=table) + + decision = planner.prepare_exact_sps_decision(num_gen_requests=0, rows=[]) + + assert decision is not None + assert decision.num_requests == 0 + assert tuple(decision.survival.shape) == (0, 5) + assert decision.native_expected_yield == 0.0 + assert decision.compact_expected_yields == (0.0, 0.0) + assert planner.allocate_exact_sps_candidate( + decision, graph_batch_size=4, verifier_budget=14 + ) == ([], 0, 3) + + +def test_exact_policy_reads_confidence_by_request_row() -> None: + table = ExactSpsCostTable( + tables={4: ExactSpsCostRow(token_counts=(0, 14), step_time_ms=(8.0, 7.0))}, + max_draft_len=5, + ) + planner = DSparkVerifyPlanner(cfg=_config(), cost_table=table) + logits = torch.full((5, 5), -8.0) + logits[3] = 8.0 + _publish_snapshot(planner, logits) + + decision = planner.prepare_exact_sps_decision(num_gen_requests=2, rows=[3, 0]) + + assert decision is not None + assert decision.survival[0].sum() > decision.survival[1].sum() + + +def test_exact_allocator_spends_the_modeled_real_row_target() -> None: + table = ExactSpsCostTable( + tables={4: ExactSpsCostRow(token_counts=(0, 22), step_time_ms=(8.0, 7.0))}, + max_draft_len=5, + ) + planner = DSparkVerifyPlanner(cfg=_config(), cost_table=table) + decision = ExactSpsLocalDecision( + num_requests=3, + survival=torch.ones((3, 5)), + native_expected_yield=18.0, + compact_expected_yields=(18.0,), + ) + + lens, budget, pad_tokens = planner.allocate_exact_sps_candidate( + decision, graph_batch_size=4, verifier_budget=22 + ) + + assert lens == [5, 5, 5] + assert budget == 12 + assert pad_tokens == 4 + assert sum(value + 1 for value in lens) + pad_tokens == 22 + + +def test_device_window_allocator_retains_ranked_host_fallback() -> None: + table = ExactSpsCostTable( + tables={3: ExactSpsCostRow(token_counts=(0, 14), step_time_ms=(8.0, 7.0))}, + max_draft_len=5, + ) + planner = DSparkVerifyPlanner(cfg=_config(), cost_table=table, device_windows=True) + decision = ExactSpsLocalDecision( + num_requests=3, + # Deliberately favor the last two requests. The host allocation remains + # a valid ranked fallback whenever the device prologue cannot run. + survival=torch.tensor([[0.1] * 5, [0.2] * 5, [0.9] * 5], dtype=torch.float32), + native_expected_yield=9.0, + compact_expected_yields=(8.0,), + ) + + lens, budget, pad_tokens = planner.allocate_exact_sps_candidate( + decision, graph_batch_size=3, verifier_budget=14 + ) + + assert lens == [1, 5, 5] + assert budget == 8 + assert pad_tokens == 0 + assert sum(value + 1 for value in lens) == 14 + + +def test_unmeasured_exact_cell_is_rejected() -> None: + table = ExactSpsCostTable( + tables={4: ExactSpsCostRow(token_counts=(0, 14), step_time_ms=(8.0, 7.0))}, + max_draft_len=5, + ) + planner = DSparkVerifyPlanner(cfg=_config(), cost_table=table) + decision = ExactSpsLocalDecision( + num_requests=2, + survival=torch.ones((2, 5)), + native_expected_yield=12.0, + compact_expected_yields=(10.0,), + ) + + with pytest.raises(ValueError, match="Unmeasured exact SPS cell"): + planner.allocate_exact_sps_candidate(decision, graph_batch_size=4, verifier_budget=16) + + +def test_policy_fails_closed_on_incomplete_snapshot_mapping() -> None: + table = ExactSpsCostTable( + tables={4: ExactSpsCostRow(token_counts=(0, 14), step_time_ms=(8.0, 7.0))}, + max_draft_len=5, + ) + planner = DSparkVerifyPlanner(cfg=_config(), cost_table=table) + _publish_snapshot(planner, torch.rand(8, 5)) + + assert planner.prepare_exact_sps_decision(num_gen_requests=4, rows=[0, 1]) is None + assert planner.stats["fallback_short_snapshot"] == 1 + + +def test_device_window_capacity_uses_the_older_snapshot() -> None: + table = ExactSpsCostTable( + tables={4: ExactSpsCostRow(token_counts=(0, 14), step_time_ms=(8.0, 7.0))}, + max_draft_len=5, + ) + planner = DSparkVerifyPlanner(cfg=_config(), cost_table=table, device_windows=True) + planner._prev_buffer = torch.full((2, 5), -8.0) + planner._prev_event = None + planner._prev_valid = True + planner._host_buffer = torch.full((2, 5), 8.0) + planner._copy_event = None + planner._snapshot_valid = True + + older = planner.prepare_exact_sps_decision(num_gen_requests=2) + planner._prev_buffer = planner._host_buffer + current = planner.prepare_exact_sps_decision(num_gen_requests=2) + + assert older is not None + assert current is not None + assert older.survival.sum() < current.survival.sum() + + +class _CopyEvent: + """CPU event double; the tests check ownership and polling, not CUDA copies.""" + + def __init__(self) -> None: + self.ready = False + + def record(self) -> None: + self.ready = False + + def query(self) -> bool: + return self.ready + + +@pytest.mark.parametrize("device_windows", [False, True]) +def test_confidence_staging_is_lazy_and_polls_before_reading( + monkeypatch: pytest.MonkeyPatch, device_windows: bool +) -> None: + events = [] + + def new_event() -> _CopyEvent: + event = _CopyEvent() + events.append(event) + return event + + monkeypatch.setattr(torch.cuda, "Event", new_event) + table = ExactSpsCostTable( + tables={4: ExactSpsCostRow(token_counts=(0, 14), step_time_ms=(8.0, 7.0))}, + max_draft_len=5, + ) + planner = DSparkVerifyPlanner(cfg=_config(), cost_table=table, device_windows=device_windows) + planner.stage_confidence(torch.empty((0, 5))) + assert events == [] + assert planner._host_buffer is None + + logits = torch.tensor([[8.0] * 5, [-8.0] * 5]) + planner.stage_confidence(logits) + assert planner.prepare_exact_sps_decision(num_gen_requests=2) is None + if device_windows: + # Rotate the first pending copy into the capacity slot. The newer + # confidence is deliberately different and must not be substituted. + planner.stage_confidence(torch.zeros_like(logits)) + assert planner.prepare_exact_sps_decision(num_gen_requests=2) is None + events[0].ready = True + + decision = planner.prepare_exact_sps_decision(num_gen_requests=2, rows=[1, 0]) + + assert decision is not None + torch.testing.assert_close(decision.survival, torch.sigmoid(logits[[1, 0]]).cumprod(dim=1)) + assert planner._host_buffer.data_ptr() != logits.data_ptr() + + +def test_device_confidence_staging_reuses_two_completed_buffers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(torch.cuda, "Event", _CopyEvent) + planner = DSparkVerifyPlanner(cfg=_config(), device_windows=True) + planner.stage_confidence(torch.full((2, 5), 1.0)) + first_buffer = planner._host_buffer + first_event = planner._copy_event + first_event.ready = True + planner.stage_confidence(torch.full((2, 5), 2.0)) + second_buffer = planner._host_buffer + planner._copy_event.ready = True + planner.stage_confidence(torch.full((2, 5), 3.0)) + + assert planner._host_buffer is first_buffer + assert planner._copy_event is first_event + assert planner._prev_buffer is second_buffer + assert planner._ready_snapshot() is None + torch.testing.assert_close(planner._ready_older_snapshot(), torch.full((2, 5), 2.0)) + torch.testing.assert_close(first_buffer, torch.full((2, 5), 3.0)) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_schedule.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_schedule.py new file mode 100644 index 000000000000..fb53b48e5ff1 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_schedule.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSpark per-request allocation tests.""" + +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark_schedule import ( + DSparkScheduleConfig, + compute_survival, + schedule_verify_lens_topk, +) + +BLOCK = 7 + + +def _cfg(**kwargs) -> DSparkScheduleConfig: + return DSparkScheduleConfig(block_size=BLOCK, **kwargs) + + +def test_allocation_is_a_prefix_and_respects_bounds(): + torch.manual_seed(3) + cfg = _cfg(min_verify_len=1) + surv = compute_survival(torch.rand(12, BLOCK)) + lens = schedule_verify_lens_topk(survival=surv, budget=20, cfg=cfg) + assert lens.dtype == torch.int32 + assert torch.all(lens >= cfg.min_verify_len) + assert torch.all(lens <= cfg.resolved_max_verify_len) + # Budget is spent above the floor and never exceeded. + assert int((lens - cfg.min_verify_len).sum()) <= 20 + # Zero budget degenerates to the floor for everyone. + cfg_floor = _cfg(min_verify_len=2) + lens = schedule_verify_lens_topk( + survival=compute_survival(torch.rand(5, BLOCK)), budget=0, cfg=cfg_floor + ) + assert torch.equal(lens, torch.full((5,), 2, dtype=torch.int32)) + # A huge budget saturates at the cap. + cfg_cap = _cfg(min_verify_len=1, max_verify_len=4) + lens = schedule_verify_lens_topk( + survival=compute_survival(torch.full((5, BLOCK), 0.99)), budget=10**6, cfg=cfg_cap + ) + assert torch.equal(lens, torch.full((5,), 4, dtype=torch.int32)) + # An empty batch allocates nothing. + assert ( + schedule_verify_lens_topk(survival=torch.zeros(0, BLOCK), budget=5, cfg=_cfg()).numel() == 0 + ) + + +def test_survival_eps_excludes_hopeless_positions(): + """Even with unlimited budget, sub-eps candidates are never admitted.""" + cfg = _cfg(min_verify_len=1, survival_eps=1e-3) + conf = torch.tensor([[1.0, 1.0, 1e-9, 1.0, 1.0, 1.0, 1.0]]) + lens = schedule_verify_lens_topk(survival=compute_survival(conf), budget=10**6, cfg=cfg) + # positions 0,1 survive; from position 2 on, survival collapses below eps. + assert int(lens[0]) == 2 + + +def test_tied_confidences_are_balanced_and_deterministic(): + """Ties form a stable balanced prefix on every rank.""" + cfg = _cfg(min_verify_len=1) + surv = compute_survival(torch.full((9, BLOCK), 0.9)) + lens = schedule_verify_lens_topk(survival=surv, budget=13, cfg=cfg) + # ``budget`` is allocated above the one-token floor already granted to + # every row, so the result contains 9 floor tokens plus 13 winners. + assert lens.tolist() == [3, 3, 3, 3, 2, 2, 2, 2, 2] + assert int(lens.sum()) == 9 + 13 + assert torch.equal(schedule_verify_lens_topk(survival=surv, budget=13, cfg=cfg), lens) + + +def test_schedule_config_rejects_bad_bounds(): + with pytest.raises(ValueError, match="min_verify_len must be >= 1"): + DSparkScheduleConfig(block_size=BLOCK, min_verify_len=0) + with pytest.raises(ValueError, match="block_size must be >= 1"): + DSparkScheduleConfig(block_size=0) + with pytest.raises(ValueError, match="< min_verify_len"): + DSparkScheduleConfig(block_size=BLOCK, min_verify_len=5, max_verify_len=3) + with pytest.raises(ValueError, match=r"\[bs, K\]"): + compute_survival(torch.rand(BLOCK)) + + +def _block5_cfg() -> DSparkScheduleConfig: + return DSparkScheduleConfig(block_size=5, min_verify_len=1, max_verify_len=5) + + +def test_topk_hands_longer_windows_to_more_confident_requests(): + """The budget must follow survival, not batch position.""" + # survival[r, k] = P(the first k+1 drafted tokens all get accepted) + survival = torch.tensor( + [ + [0.99, 0.98, 0.97, 0.96, 0.95], # confident + [0.90, 0.70, 0.50, 0.30, 0.10], # middling + [0.20, 0.04, 0.01, 0.00, 0.00], # collapses + ] + ) + lens = schedule_verify_lens_topk(survival=survival, budget=6, cfg=_block5_cfg()).tolist() + + assert len(lens) == 3 + assert lens[0] > lens[2], ( + f"confident request got {lens[0]}, collapsing request got {lens[2]}; " + f"the budget must follow survival, not batch position" + ) + assert lens[0] >= lens[1] >= lens[2] + # Every request keeps at least the floor, and the budget is respected. + assert min(lens) >= 1 + assert sum(lens) - len(lens) * 1 <= 6 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 new file mode 100644 index 000000000000..5cc980a1f603 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_workload_admission.py @@ -0,0 +1,322 @@ +# 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 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, + ) + + +@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, + ) From d1b3d3ac07b16dbe46e6fb8ee03985fec6471c83 Mon Sep 17 00:00:00 2001 From: Zhaoyuan He Date: Wed, 7 Oct 2026 09:36:52 -0700 Subject: [PATCH 2/2] [TRTLLM-14620][fix] Authenticate and fully validate confidence artifacts Signed-off-by: Zhaoyuan He --- .../_torch/speculative/dspark_planner.py | 38 +++++---- .../test_dspark_workload_admission.py | 84 +++++++++++++++++++ 2 files changed, 106 insertions(+), 16 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dspark_planner.py b/tensorrt_llm/_torch/speculative/dspark_planner.py index 931dc1e94554..812a6932b2db 100644 --- a/tensorrt_llm/_torch/speculative/dspark_planner.py +++ b/tensorrt_llm/_torch/speculative/dspark_planner.py @@ -372,7 +372,12 @@ def load_runtime_sps_cost_table( 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.""" + """Load an authenticated schema-v2 exact ``T(G,V)`` table. + + 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. + """ payload = _read_sps_cost_payload(path) schema_version = payload.get("schema_version") if type(schema_version) is not int or schema_version != 2: @@ -961,15 +966,12 @@ def evaluate_confidence_workload_admission( } -def _sha256_file(path: str | Path) -> str: - digest = hashlib.sha256() - with Path(path).open("rb") as file: - for chunk in iter(lambda: file.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() +def _read_pinned_bytes(path: str | Path) -> tuple[bytes, str]: + data = Path(path).read_bytes() + return data, hashlib.sha256(data).hexdigest() -def _read_strict_json_object(path: str | Path, *, name: str) -> dict[str, object]: +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: @@ -978,8 +980,7 @@ def reject_duplicate_fields(pairs: list[tuple[str, object]]) -> dict[str, object parsed[key] = value return parsed - with Path(path).open(encoding="utf-8") as file: - payload = json.load(file, object_pairs_hook=reject_duplicate_fields) + 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 @@ -1005,11 +1006,11 @@ def load_confidence_workload_admission( expected_receipt_sha256, field="expected confidence workload admission receipt SHA256", ) - receipt_sha256 = _sha256_file(path) + 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( - _read_strict_json_object(path, name="confidence workload admission receipt"), + _parse_strict_json_object(receipt_bytes, name="confidence workload admission receipt"), name="confidence workload admission receipt", fields=_WORKLOAD_ADMISSION_ENVELOPE_FIELDS, ) @@ -1049,15 +1050,19 @@ def load_confidence_workload_admission( _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") - if admission["sps_cost_table_sha256"] != _sha256_file(sps_cost_table_path): + 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") - if admission["live_engine_fingerprint_sha256"] != _sha256_file(live_engine_fingerprint_path): + 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( - _read_strict_json_object(live_engine_fingerprint_path, name="live engine fingerprint"), + _parse_strict_json_object(live_fingerprint_bytes, name="live engine fingerprint"), name="live engine fingerprint", fields=_V2_FINGERPRINT_FIELDS, ) @@ -1076,13 +1081,14 @@ def load_confidence_workload_admission( 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 = _read_strict_json_object(sps_cost_table_path, name="SPS cost artifact") + 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"], 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 index 5cc980a1f603..f8026e353511 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_workload_admission.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_workload_admission.py @@ -4,6 +4,7 @@ import hashlib import json +from copy import deepcopy from pathlib import Path import pytest @@ -232,6 +233,89 @@ def test_workload_admission_rejects_changed_artifact_bytes(tmp_path: Path, artif ) +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)