diff --git a/examples/models/core/deepseek_v4/README.md b/examples/models/core/deepseek_v4/README.md index 5672ef03a73e..ddac8fa83a45 100644 --- a/examples/models/core/deepseek_v4/README.md +++ b/examples/models/core/deepseek_v4/README.md @@ -472,6 +472,211 @@ MXFP4 routed MoE expert weights, TensorRT LLM automatically applies the routed-e configuration. +## DSpark confidence verification (prototype) + +Confidence-guided verification is **default-off and experimental**. It proposes the full physical +draft block `K`, then may select a smaller, measured verifier-token budget. It does not reduce the +physical drafter length. This draft integration is not a production deployment recipe: strict +receipt generation and performance qualification are external, and the composed source still needs native +build, GPU/full-model, accuracy and performance validation. Optional mixed scheduling has the +separate evidence and route requirements below; no composed-tip GPU qualification is implied. + +`K` is the physical draft length, `G` is the rank-local request count for a captured graph, and +`V` is its verifier-token budget. In the cost table, `V=0` denotes the native full-width path; +positive `V` values denote measured compact layouts. The table must cover the configured physical +`K` and exact graph ladder. The scheduler never invents a missing cell. + +### Requested configuration versus actual execution + +Set `DSparkDecodingConfig.enable_confidence_scheduling=True` to opt in. By default, structurally +valid existing costs and calibration can be reused even when their recorded source or GPU-node +metadata differ from the current run. A warning makes clear that predictions may be stale and +that positive performance economics have not been established. Reuse does not change target +verification or the acceptance rule. + +Set `validate_confidence_artifact_freshness=True` to require the previous strict live-fingerprint +and positive workload-admission checks before constructing engines and workers: + +| Startup result | Resolved execution | +| --- | --- | +| Explicit opt-in with default artifact reuse | Confidence stays enabled; measured costs must still pass schema, finiteness, physical-`K`, exact graph and layout checks. No admission receipt is created or asserted. | +| Strict validation with missing, unreadable, stale or non-positive admission evidence | Both confidence flags become false and all confidence artifact paths/digests are cleared; ordinary feature-off static `K` runs. A warning explains the fallback. | +| Physical `K=1` | Feature-off static `K1`; there is no smaller verifier tier. Confidence inputs are ignored and cleared. | +| Strict validation with a matching receipt and positive aggregate value | Confidence remains requested; the existing environment, graph, checkpoint and runtime checks must still pass. | + +An **enabled-native** policy decision is different from startup +feature-off fallback: it retains the enabled control path while verifying the native layout. +An enabled compact decision uses a positive measured budget. Optional fused-scheduler failure +falls back to the tensor scheduler for the affected graph size; that is a third, separate fallback. + +Before reporting an "enabled" experiment, inspect the resolved +`TorchLlmArgs.speculative_config.enable_confidence_scheduling` and record both the requested and +resolved configuration. Require observed native/compact routing and completed outputs, rather +than inferring execution from an input YAML flag or a successful server start. Do not count +startup feature-off fallback as an enabled-native run. With confidence disabled, omit all +confidence-only fields; explicitly supplying them is a configuration error. + +Confidence scheduling requires padded CUDA graphs and PP=1. Guided decoding, LoRA and incompatible +sparse top-k settings remain unsupported. Per-request logits post-processors are rejected before +activation. Advanced rejection draft sampling is also unsupported: its sampled proposal prefix +differs from the greedy prefix that produced the confidence scores. Greedy requests (including +`top_k=1`) and requests with rejection draft sampling disabled retain their existing behavior. +Confidence-disabled execution retains its ordinary sampling and post-processing paths. + +### Reusing costs and calibration + +For default reuse, supply `confidence_sps_table_path` and, when using a calibrated head, +`confidence_sts_path`. Costs must be finite positive schema-v2 measurements with native and compact +cells for every configured graph size; STS temperatures must be a finite positive vector with +one value per physical draft position. Invalid artifacts raise errors rather than being silently +timed as enabled static execution. Their original metadata is retained; the runner records its +observed runtime separately instead of rewriting the artifact's provenance. + +For example, the confidence portion of a padded-graph configuration is: + +```yaml +speculative_config: + decoding_type: DSpark + max_draft_len: 3 + enable_confidence_scheduling: true + confidence_sps_table_path: /your/artifacts/k3-costs.json + confidence_sts_path: /your/artifacts/k3-sts.json +``` + +This is a configuration example, not an exact-tip GPU-validated reproduction recipe. Keep the +same checkpoint, physical `K`, inputs, parallelism, graph ladder and other runtime options when +comparing it with static DSpark (`enable_confidence_scheduling: false`, without confidence-only +fields). A fair measured comparison keeps source, configuration, inputs and output work matched, +uses independent model lifetimes, checks actual routing and completed outputs, and records actual +timings. Reused predictions do not invalidate that measurement: retain their original provenance +and label freshness unchecked, without promising the earlier gain. Fresh profiling is optional +and can improve prediction quality. Changing batch size requires a graph ladder covered by the +supplied table, not a separate per-batch-size package. + +Both modes retain the table's immediate predicted-goodput margin (normally 1%). If the table +supplies an `iteration_drain_guard`, its prices and final-drain policy are unchanged. Without that +optional metadata, default reuse selects by immediate goodput alone; strict validation retains +native verification. Neither behavior certifies an end-to-end gain. + +### Optional freshness and workload-admission validation + +With `validate_confidence_artifact_freshness: true`, provide all four prototype inputs: + +- `confidence_sps_table_path`: schema-v2 measured costs with exact physical `K`, rank-local graph + sizes `G`, native `V=0` cells and supported positive verifier budgets. +- `confidence_sps_live_fingerprint_path`: an independently supplied fingerprint for that runtime + and graph ladder, not a copy relabeled from an older cost artifact. +- `confidence_admission_receipt_path`: the schema-v1 workload receipt. +- `confidence_admission_receipt_sha256`: the expected lowercase SHA256 of the **receipt file + bytes**, obtained from a trusted, separately reviewed deployment package. + +There is no supported public cost-profile/calibration/receipt-generation command in this source +change. Self-generation requires the separate profiling tooling and its validated instructions; +the fields above are consumers, not a cost collector. Do not manufacture costs or an admitting +receipt from unit-test fixtures. A strict receipt's producer must preserve the measured cost, +workload, calibration, selector-replay and runtime evidence and review their applicability to the +actual deployment. A source, calibration, physical-`K` or workload change requires reassessment; +historical scores or timing results cannot be relabeled as new-source qualification. + +The receipt envelope contains exactly `schema_version: 1`, `admission` and `admission_sha256`. +The last field hashes the canonical admission-body JSON +(`json.dumps(body, sort_keys=True, separators=(",", ":")).encode()`). +The body contains: + +- `physical_k`, `source_head`, `source_diff_sha256`, `runtime_snapshot`; +- `sps_cost_table_sha256`, `live_engine_fingerprint_sha256`, `selector_identity_sha256`; +- `workload_identity_sha256`, `calibration_result_sha256`, `selector_replay_sha256`; +- `policy_steps`, `compact_choices`, `gross_compact_value_ms_lower_bound`, + `fixed_confidence_path_overhead_ms_upper_bound`, `safety_margin_ms`, `admitted`. + +The selector identity must match `EXACT_SPS_SELECTOR_IDENTITY_SHA256` in +`tensorrt_llm/_torch/speculative/dspark_planner.py`. Replay the existing marginal selector without +allocating fixed overhead back to candidates through a `policy_steps / compact_choices` ratio. +Admission recomputes: + +```text +net_value_lower_bound_ms = + gross_compact_value_ms_lower_bound + - policy_steps * fixed_confidence_path_overhead_ms_upper_bound + - safety_margin_ms + +admitted = compact_choices > 0 and net_value_lower_bound_ms > 0 +``` + +Native decisions still pay the enabled confidence-path overhead, so **all** policy steps count. +Strict admission does not replace per-decision selection: exact measured tiers and the existing +iteration/drain guard still determine whether a compact choice is allowed. + +Hashes provide integrity and cross-artifact consistency checks, **not signatures or proof of +origin**. The runtime checks supplied receipt, cost and fingerprint values; it does not independently +derive the running checkpoint/workload identity, re-execute the calibration or selector replay, +or verify their referenced result files. The external producer/deployment package remains the +trust boundary. Hashing an arbitrary receipt and passing that same digest does not establish +trust or qualify an enabled experiment. + +Contract regressions are in `test_dspark_workload_admission.py`, +`test_dspark_admission_resolution.py` (both under +`tests/unittest/_torch/speculative/hw_agnostic/`) and +`tests/unittest/llmapi/test_dspark_workload_admission_config.py`. These fixtures are synthetic +validator tests, not deployment evidence. `test_dspark_artifact_reuse.py` covers optional +freshness, unchanged structural guards and the absent-drain behavior. + +### Optional mixed context/generation scheduling + +`confidence_mixed_evidence` is an optional nested object. Its default is `None`, which leaves +mixed steps at native verification width and does not change enabled decode-only scheduling. +Mixed costs are not decode `T(G,V)` costs: context chunks and their positions, generation history, +bootstrap requests, ordered ranks and the shared body capacity all affect the mixed cost key. + +The object contains `costs_path`, `costs_sha256`, `admission_path`, `admission_sha256`, +`protocol_sha256`, `live_fingerprint_path` and `live_fingerprint_sha256`. Cost, receipt and live +fingerprint digests pin exact file bytes; the protocol digest identifies the measured protocol. +These pins come from a trusted deployment package. The live fingerprint must be observed independently +for the current source, runtime, model, GPU class, topology and selector, not copied from the cost +table. The cost file uses `dspark-mixed-breakable-costs-v1`; its separate receipt uses +`dspark-mixed-breakable-admission-v1`. The reuse setting applies to decode schema-v2 inputs only; +it does not relax this separate mixed-execution contract. + +The supported mixed route is embedded DSpark with physical `K=3`, TP8, PP1, CP1, attention DP, +greedy sampling and captured BREAKABLE body buckets `[96, 128, 192, 256, 384, 512]`. Standalone +draft models, unsupported routes and uncovered mixed geometry retain native mixed verification. +Missing or invalid mixed evidence does not disable an already enabled decode policy. +First context chunks also remain native because KV preparation may still resolve a prefix-cache +hit and change their geometry. Later chunks are eligible only when their full-width native rows +fit the configured token budget, so resource preparation does not trim a quoted context shape. + +When a peer rank prefills, all ranks use the existing shared BREAKABLE padding path rather than +replaying a rank-local full decode graph. Only authenticated generation owners may receive shorter +verification prefixes; context rows, bootstrap requests and graph-padding requests retain their +required rows. The runner publishes those lengths as host-authoritative before attention metadata +preparation, without changing the attention backend, sampler or acceptance rule. Stale or unready +confidence owners retain full physical `K`. + +Selection requires an exact measured geometry/action cell and retains the minimum 1% predicted +goodput gate. The current pricing field is `rank0_forward_pricing_proxy`, not a measurement of the +whole distributed step. Workload admission must charge the opt-in path's overhead on **all** policy +steps, including all-decode and native choices: the optional path adds CPU curve preparation, +producer-stamp transfer and a fixed trailer to the existing policy allgather. Policy equivalence +alone is not evidence of zero overhead or preserved throughput. + +There is no public mixed receipt-generation command in this change. Supplying valid-shaped JSON +or hashing test fixtures does not qualify a deployment. A separately reviewed producer must supply +matching measured costs, runtime/accuracy results, selector replay and positive workload economics. +Exact-tip GPU/full-model validation and production mixed evidence remain required before making a +deployment or performance claim. The CPU contracts are in `test_dspark_mixed_evidence.py` and +`test_dspark_mixed_runtime.py` under `tests/unittest/_torch/speculative/hw_agnostic/`. + +### Drafter fusion scope + +The BF16 CuTe attention and RMSNorm–RoPE path is eligible for physical `K=1..8` under the existing +SM100/SM103, geometry and layout guards; unsupported inputs retain the reference path. This +dispatch is independent of the confidence flag, so it also affects static DSpark. The eight-row +storage bound is unchanged; no arbitrary-`K` kernel extension is claimed. + +The source-matched extension passed 98 GPU unit cases on B300/SM103. That is not full-model or +performance qualification of this composed integration. SM100 and exact-integration GPU/model +validation remain pending; keep this scope experimental until those gates and review complete. + ## Notes and Troubleshooting - `DeepSeek-V4 on Hopper requires kv_cache_config.dtype='fp8_ds_mla'`: remove an explicit BF16/FP8 diff --git a/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py index d414a99f60a8..8feab7b891bb 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py @@ -5,7 +5,7 @@ The op runs the tcgen05 MMA kernel (:class:`DSparkAttention`) on the supported DSV4 DSpark -geometry: 128 heads, head_dim 512, draft block 5 or 6, and a 128-row rolling +geometry: 128 heads, head_dim 512, draft block 1 through 8, and a 128-row rolling window. Other shapes fall back to the pure-PyTorch reference path in ``models/modeling_dspark.py``. """ @@ -32,7 +32,8 @@ _DSV4_DSPARK_HEAD_DIM = 512 _DSV4_DSPARK_WINDOW_SIZE = 128 _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE = 8 -_DSV4_DSPARK_BLOCK_SIZES = (5, 6) +# Physical K is bounded by the eight-row draft-buffer ABI, not the MMA tile. +_DSV4_DSPARK_BLOCK_SIZES = (1, 2, 3, 4, 5, 6, 7, 8) _DSV4_DSPARK_ROPE_DIM = 64 @@ -92,7 +93,9 @@ def is_fused_dsv4_dspark_attention_supported( f"expected q/kv_cache ranks 4/3, got {q.ndim}/{kv_cache.ndim}", "tensor_ranks" ) if q.shape[1] not in _DSV4_DSPARK_BLOCK_SIZES: - return _log_unsupported(f"draft block size must be 5 or 6, got {q.shape[1]}", "block_size") + return _log_unsupported( + f"draft block size must be 1 through 8, got {q.shape[1]}", "block_size" + ) if q.shape[2:] != (_DSV4_DSPARK_NUM_HEADS, _DSV4_DSPARK_HEAD_DIM): return _log_unsupported( "q must have 128 heads and head_dim 512; " @@ -346,7 +349,7 @@ def fused_dsv4_dspark_attention( ): raise ValueError( "fused_dsv4_dspark_attention requires contiguous supported DSV4 DSpark tensors " - "([B, 5|6, 128, 512] queries, [B, 8, 512] draft blocks, INT32 indices, " + "([B, K, 128, 512] queries with 1 <= K <= 8, [B, 8, 512] draft blocks, INT32 indices, " "and a 128-row BF16 window) on SM100 or SM103" ) diff --git a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py index 56bc9f62d707..f2e0b741b2b5 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py @@ -20,7 +20,7 @@ _DSV4_DSPARK_ROPE_DIM = 64 _DSV4_DSPARK_WINDOW_SIZE = 128 _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE = 8 -_DSV4_DSPARK_BLOCK_SIZES = (5, 6) +_DSV4_DSPARK_BLOCK_SIZES = (1, 2, 3, 4, 5, 6, 7, 8) _DSV4_DSPARK_ARCH_BY_SM = { 100: "sm_100", 103: "sm_103", diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark/attention.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark/attention.py index 458d2877f12d..098458969918 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark/attention.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark/attention.py @@ -195,7 +195,7 @@ def __new_from_mlir_values__(self, values): class DSparkAttention(DSparkAttentionKernel): - """Attention over a 128-token rolling window and one 5/6-token draft block.""" + """Attention over a 128-token rolling window and one supported draft block.""" window_size = 128 block_size = 6 @@ -260,8 +260,9 @@ def __init__( ] if mismatches: raise ValueError("Unsupported DSpark kernel configuration: " + ", ".join(mismatches)) - if seq_len_q not in (5, 6): - raise ValueError(f"DSpark block size must be 5 or 6, got {seq_len_q}") + # Each query gets its own head tile; only draft storage bounds K. + if seq_len_q not in (1, 2, 3, 4, 5, 6, 7, 8): + raise ValueError(f"DSpark block size must be 1 through 8, got {seq_len_q}") if inverse_rope_dim not in (0, 64): raise ValueError(f"DSpark inverse_rope_dim must be 0 or 64, got {inverse_rope_dim}") diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 030b34272450..262406affb23 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -2103,13 +2103,23 @@ def write_context_windows_batched(self, main_hidden, positions, slots, mask, kv_ ) def _prepare_confidence_head_for_load(self, weights: Dict) -> None: - """Match the projection's optional bias to the checkpoint before loading.""" + """Prepare and validate the embedded confidence namespace before loading.""" head = self.dspark_model.mtp_layers[-1].confidence_head - bias_key = f"mtp_layers.{self.num_stages - 1}.confidence_head.proj.bias" - if head is not None and bias_key in weights and head.proj.bias is None: + prefix = f"mtp_layers.{self.num_stages - 1}.confidence_head." + module_weights = { + name[len(prefix) :]: value for name, value in weights.items() if name.startswith(prefix) + } + if head is None: + if module_weights: + raise ValueError("DSpark confidence head checkpoint has no constructed head") + return + if "proj.bias" in module_weights and head.proj.bias is None: # This runs before weight loading and graph capture. Empty storage # avoids consuming RNG or reinitializing the existing projection. head.proj.bias = nn.Parameter(head.proj.weight.new_empty(head.proj.out_features)) + # The generic V4 loader skips this parameter-free head container and + # visits its Linear child, so strict validation must run at this seam. + head._validate_checkpoint_weights(module_weights) def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): """Load the ``mtp.*`` draft weights from the (full) checkpoint dict. diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 569ea56dced9..52d3134b55d4 100644 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -579,6 +579,13 @@ def load_sts_temperatures(self, temperatures: torch.Tensor) -> None: def load_weights(self, weights: list) -> None: """Load every confidence-head checkpoint parameter strictly.""" (module_weights, ) = weights + self._validate_checkpoint_weights(module_weights) + for name, param in self.named_parameters(): + param.data.copy_(module_weights[name][:]) + + def _validate_checkpoint_weights( + self, module_weights: dict[str, torch.Tensor]) -> None: + """Validate a head namespace before generic embedded parameter loading.""" expected = dict(self.named_parameters()) unexpected = [key for key in module_weights if key not in expected] missing = [key for key in expected if key not in module_weights] @@ -593,7 +600,11 @@ def load_weights(self, weights: list) -> None: f"DSpark confidence head checkpoint is missing key(s) {sorted(missing)}" ) for name, param in expected.items(): - param.data.copy_(module_weights[name][:]) + actual_shape = tuple(module_weights[name].shape) + if actual_shape != tuple(param.shape): + raise ValueError( + f"DSpark confidence head checkpoint {name} has shape {actual_shape}; " + f"expected {tuple(param.shape)}") class Eagle3Attention(Attention): diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 7aedaf8a81e6..2a725292ced4 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -83,12 +83,23 @@ from .seq_slot_manager import SeqSlotManager GB = 1 << 30 +_DEFAULT_INDEX_MAPPER_RESERVED_SLOTS = 1 +_CONFIDENCE_INDEX_MAPPER_RESERVED_SLOTS = 2 def ceil_div(a: int, b: int) -> int: return (a + b - 1) // b +def _get_num_reserved_index_slots(spec_config: Optional[SpeculativeConfig], + estimating_kv_cache: bool) -> int: + """Return IndexMapper headroom required by the execution policy.""" + if getattr(spec_config, "enable_confidence_scheduling", False): + return (0 if estimating_kv_cache else + _CONFIDENCE_INDEX_MAPPER_RESERVED_SLOTS) + return _DEFAULT_INDEX_MAPPER_RESERVED_SLOTS + + def _get_initial_lora_data_type( configured_lora_data_type: Optional[torch.dtype], ) -> Optional[torch.dtype]: @@ -3246,6 +3257,8 @@ def _create_kv_cache_manager( manager_extra_kwargs["max_copy_beam_width"] = max_copy_beam_width manager_extra_kwargs[ "cold_page_codec_provider"] = cold_page_codec_provider + manager_extra_kwargs["num_reserved_index_slots"] = ( + _get_num_reserved_index_slots(spec_config, estimating_kv_cache)) manager_extra_kwargs["kv_events_config"] = kv_events_config manager_extra_kwargs["joint_kv_cache_reuse"] = joint_kv_cache_reuse manager_extra_kwargs[ diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 19ed990f33a1..3555b4204efc 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -28,8 +28,8 @@ from ..speculative.utils import get_draft_kv_cache_manager from ..utils import make_weak_ref, piecewise_cuda_graph from .llm_request import LlmRequest, get_draft_token_length -from .resource_manager import (BaseResourceManager, ResourceManager, - ResourceManagerType) +from .resource_manager import (BaseResourceManager, NoFreeSlotsError, + ResourceManager, ResourceManagerType) from .sampler import SampleStateTensors from .sampler.sampler_common import SampleType from .scheduler import ScheduledRequests @@ -46,6 +46,32 @@ def _strict_buffer_check_enabled() -> bool: # A large prime number used for dummy request IDs to avoid collisions CUDA_GRAPH_DUMMY_REQUEST_ID = (1 << 64) - 1 + + +def cuda_graph_dummy_request_id(runtime_draft_len: int, *, variant: int, + max_draft_len: int) -> int: + """Return the disjoint padding-dummy ID for a draft/window variant.""" + runtime_draft_len = int(runtime_draft_len) + variant = int(variant) + max_draft_len = int(max_draft_len) + if (variant not in (0, 1) or max_draft_len < 0 or runtime_draft_len < 0 + or runtime_draft_len > max_draft_len): + raise ValueError("invalid CUDA padding dummy identity: " + f"draft_len={runtime_draft_len}, variant={variant}, " + f"max_draft_len={max_draft_len}") + return (CUDA_GRAPH_DUMMY_REQUEST_ID - variant * (max_draft_len + 1) - + runtime_draft_len) + + +def _padding_dummy_draft_token_capacity( + spec_config: Optional[DecodingBaseConfig], + allocated_draft_tokens: int) -> int: + """Keep DSpark's physical-K bound without restricting tree speculation.""" + if getattr(spec_config, "decoding_type", None) == "DSpark": + return int(spec_config.max_draft_len) + return int(allocated_draft_tokens) + + # Gen dummies get prompt_len = token_num - 1. Before capturing enc-dec decode # graphs, EncoderDecoderRunner._prepare_capture_batch temporarily runs each # dummy generation request as a one-token context chunk to write its cross-KV @@ -58,6 +84,45 @@ def get_mrope_dummy_seq_slot(max_num_tokens: int, pp_size: int) -> int: return max_num_tokens * pp_size +@dataclass(frozen=True) +class ADPShapeAgreement: + """One DSpark policy collective's queue and graph-shape agreement. + + The object is local process state, so ``batch_identity`` is intentionally + the Python identity on this rank. The collective-derived fields are equal + across ranks. Only a finalized exact-policy agreement may suppress the + later graph-shape and padding collectives. + """ + + iteration: int + batch_identity: int + local_batch_size: int + peer_batch_sizes: Tuple[int, ...] + all_can_graph: bool + widest_batch_size: int + graph_batch_size: int + draft_len: int + padding_ready: bool + ragged_bucket: Optional[int] = None + reuse_graph_shape: bool = False + + @property + def can_queue(self) -> bool: + return bool(self.peer_batch_sizes) and all(self.peer_batch_sizes) + + @property + def can_queue_this_rank(self) -> bool: + return self.local_batch_size > 0 + + def matches(self, batch: ScheduledRequests, iteration: int, *, + padded: bool) -> bool: + expected_size = (self.graph_batch_size + if padded else self.local_batch_size) + return (self.iteration == int(iteration) + and self.batch_identity == id(batch) + and batch.batch_size == expected_size) + + class KeyType(NamedTuple): batch_size: int draft_len: int @@ -76,6 +141,9 @@ class KeyType(NamedTuple): num_encoder_tokens: int = 0 peft_cache_data_type: Optional[torch.dtype] = None use_lora_graph: bool = False + # Total submitted verifier tokens for a ragged generation graph. ``None`` + # retains the existing uniform graph family. + ragged_verify_bucket: Optional[int] = None def _save_spec_decode_capture_state( @@ -173,6 +241,12 @@ def __init__(self, config: CUDAGraphRunnerConfig): self.max_supported_batch_size = config.max_cuda_graph_batch_size self.max_beam_width = config.max_beam_width self.spec_config = config.spec_config + self._dspark_confidence_enabled = bool( + self.spec_config is not None and getattr( + self.spec_config, "enable_confidence_scheduling", False)) + self._dspark_trims_submitted_tokens = bool( + self.spec_config is not None and getattr( + self.spec_config, "enable_confidence_scheduling", False)) self.sparse_config = config.sparse_attention_config self.is_encoder_decoder = config.is_encoder_decoder self.enable_encoder_decoder_mixed_cuda_graph = ( @@ -191,6 +265,16 @@ def __init__(self, config: CUDAGraphRunnerConfig): self._pre_replay_hooks: List[Callable[ [KeyType, int, Dict[str, Any], Dict[str, Any]], None]] = [] self._pre_replay_position_ids: Optional[torch.Tensor] = None + # Low window carried by ordinary ragged padding rows. Zero-real exact + # cells may use a second cached dummy at low+1 to represent V=G*q+r. + self.ragged_pad_verify_len = 0 + # Number of high-token rows in the complete padded graph, including the + # already scheduled attention-DP dummy. Zero is the resting value. + self.ragged_zero_real_high_rows = 0 + # Published by ModelEngine when the ragged fit picks a token bucket; + # None is the resting value (no ragged step in flight). + self.agreed_ragged_bucket: Optional[int] = None + self.adp_shape_agreement: Optional[ADPShapeAgreement] = None self.graphs: Dict[KeyType, torch.cuda.CUDAGraph] = {} self.graph_outputs: Dict[KeyType, @@ -198,6 +282,7 @@ def __init__(self, config: CUDAGraphRunnerConfig): self.graph_metadata: Dict[KeyType, Dict[str, Any]] = {} self.memory_pool = config.cuda_graph_mem_pool self.padding_dummy_requests: Dict[int, LlmRequest] = {} + self.secondary_padding_dummy_requests: Dict[int, LlmRequest] = {} self.dynamic_draft_len_mapping = config.dynamic_draft_len_mapping self.shared_static_tensors: Dict[str, torch.Tensor] = {} @@ -281,6 +366,11 @@ def _is_mixed_encoder_decoder_batch(self, batch: ScheduledRequests) -> bool: and batch.num_generation_requests > 0) def _can_run_cuda_graph_batch(self, batch: ScheduledRequests) -> bool: + # A peer may be prefilling while this rank contains only generation. + # Host-authoritative windows use the common BREAKABLE body, never a + # rank-local full-generation graph or its device layout prologue. + if getattr(self, "_dspark_host_window_batch", None) is batch: + return False return (batch.can_run_cuda_graph or (self.enable_encoder_decoder_mixed_cuda_graph and self._is_mixed_encoder_decoder_batch(batch))) @@ -430,7 +520,8 @@ def get_graph_key( context_query_len=context_query_len, num_encoder_tokens=num_encoder_tokens, peft_cache_data_type=peft_cache_data_type, - use_lora_graph=use_lora_graph) + use_lora_graph=use_lora_graph, + ragged_verify_bucket=self._ragged_verify_bucket(batch)) def set_capture_sample_type(self, sample_type: Optional[SampleType]) -> None: @@ -507,6 +598,39 @@ def _get_compatible_mixed_encoder_decoder_key(self, return min(compatible_keys, key=lambda captured_key: captured_key.num_encoder_tokens) + def _ragged_verify_bucket(self, batch: ScheduledRequests) -> Optional[int]: + """Total verified token count when the batch uses ragged windows. + + None unless *every* generation request carries a window; None keeps + the original (uniform) key. Config-gated before touching the batch: + this runs on every graph-key build for every model. + """ + if not self._dspark_trims_submitted_tokens: + return None + # Windows can vanish between the fit and key derivation; a ragged key + # over a window-less batch replays stale token-major row maps (IMA). + if any( + getattr(request, "py_verify_len", None) is None + for request in batch.generation_requests): + return None + # Return the fitted bucket, never a fresh sum over the batch: every + # attention-DP rank must key on the same value, and the fit's + # allgathered arithmetic agrees by construction where a local sum only + # agrees by accident. + return self.agreed_ragged_bucket + + @staticmethod + def _local_draft_len(batch: ScheduledRequests) -> int: + """This rank's draft length, for the attention-DP shape comparison. + + Tolerant where :meth:`get_graph_key` asserts: a malformed batch + reports its max, differs from well-formed peers, and the gate + declines the graph. + """ + return max((len(request.py_draft_tokens) + for request in batch.generation_requests), + default=0) + @staticmethod def _get_mrope_position_delta(request: Any) -> Optional[Any]: mrope_position_delta = getattr(request, "py_mrope_position_delta", None) @@ -568,28 +692,69 @@ def maybe_get_cuda_graph( is_mixed_encoder_decoder = self._is_mixed_encoder_decoder_batch(batch) can_run_cuda_graph = self._can_run_cuda_graph_batch(batch) batch_size = batch.batch_size - # The sampling tier joins the graph key, so it has to agree across the - # attention-DP ranks or they would replay different graphs -- the same - # reason can_run_cuda_graph and batch_size are all-gathered here. Ranks - # see different batches (an idle one resolves FULL), so fall back to - # FULL for the whole group unless every rank picked the same tier. forced_sample_type: Optional[SampleType] = None - if self.enabled and self.config.enable_attention_dp and self.config.mapping.tp_size > 1: - local_sample_type = self._resolve_sample_type( - batch, promoted_context_request_ids) + agreement = (self.adp_shape_agreement + if self._dspark_confidence_enabled else None) + # The policy agreement covers shape, not sampling tier. When main's + # in-graph sampler is enabled, retain its per-rank tier agreement. + reuse_agreement = bool( + agreement is not None and agreement.reuse_graph_shape + and not self.enable_in_graph_sampling + and agreement.matches(batch, agreement.iteration, padded=True)) + if reuse_agreement: + local_bucket = self._ragged_verify_bucket(batch) + local_draft_len = self._local_draft_len(batch) + if (not can_run_cuda_graph + or batch_size != agreement.graph_batch_size + or local_bucket != agreement.ragged_bucket + or local_draft_len != agreement.draft_len): + raise RuntimeError( + "DSpark cached ADP graph-shape agreement no longer " + "matches the local replay key") + adp_graph_agreement = (self.enabled and self.config.enable_attention_dp + and self.config.mapping.tp_size > 1) + local_sample_type = self._resolve_sample_type( + batch, + promoted_context_request_ids) if adp_graph_agreement else None + if (adp_graph_agreement and not self._dspark_confidence_enabled): + # Keep main's sampling-tier exchange unchanged when confidence is + # off; there is no ragged shape to reconcile on this path. graph_batch_info = self._gather_adp_graph_batch_info( batch, can_run_cuda_graph, local_sample_type) - all_can_run_cuda_graph = all(rank_info[0] - for rank_info in graph_batch_info) + all_can_graph = all(rank_info[0] for rank_info in graph_batch_info) all_batch_sizes_equal = all(rank_info[1] == graph_batch_info[0][1] for rank_info in graph_batch_info) - - if not all_can_run_cuda_graph or not all_batch_sizes_equal: + if not all_can_graph or not all_batch_sizes_equal: return None, None, None - forced_sample_type = (local_sample_type if all( rank_info[2] == graph_batch_info[0][2] for rank_info in graph_batch_info) else SampleType.FULL) + elif (adp_graph_agreement and not reuse_agreement): + # Compare every key component a rank can decide for itself: replay + # freezes all_rank_num_tokens, so a shape mismatch is invisible at + # replay time. draft_len is settable rank-locally (acceptance-rate + # gate), so it is compared rather than assumed. -1 = "not + # applicable" so the payload is plain ints on every rank. + local_bucket = self._ragged_verify_bucket( + batch) if can_run_cuda_graph else None + all_can_graph_batch = self.config.dist.tp_allgather_int64([ + can_run_cuda_graph, + batch_size, + -1 if local_bucket is None else int(local_bucket), + self._local_draft_len(batch) if can_run_cuda_graph else -1, + local_sample_type.value, + ]).tolist() + is_all_gen_only = all(all_can_graph[0] + for all_can_graph in all_can_graph_batch) + all_shapes_equal = all(peer[1:4] == all_can_graph_batch[0][1:4] + for peer in all_can_graph_batch) + + if not is_all_gen_only or not all_shapes_equal: + return None, None, None + + forced_sample_type = (local_sample_type if all( + peer[4] == all_can_graph_batch[0][4] + for peer in all_can_graph_batch) else SampleType.FULL) if not self.enabled or not can_run_cuda_graph: return None, None, None @@ -678,6 +843,12 @@ def get_graph_pool(self): return self.memory_pool def _get_num_tokens_for_key(self, key: KeyType) -> int: + if key.ragged_verify_bucket is not None: + if key.num_contexts: + raise ValueError( + "Ragged verifier graphs do not support mixed context batches." + ) + return int(key.ragged_verify_bucket) token_per_generation = key.draft_len + 1 return (key.num_contexts * key.context_query_len + (key.batch_size * self.max_beam_width - key.num_contexts) * @@ -989,6 +1160,25 @@ def _check_strict_buffer_stability(self, stored_meta: Dict[str, self._assert_graph_tensor_ptrs_stable( stored_meta["spec_metadata"], stored_meta["spec_metadata_ptrs"]) + def will_pad_to(self, padded_bs: int, num_requests: int) -> bool: + """Whether a batch of ``num_requests`` can actually reach ``padded_bs``. + + Mirrors the SIZE guards in :meth:`_get_padded_batch` only; the + resource-dependent bails (dummy KV capacity, encoder-decoder state) + are deliberately not predicted. The ragged fit sizes its bucket grid + from the padded row count, so a wrong yes strands the fitted total in + no captured bucket. + """ + if padded_bs <= 0 or num_requests <= 0: + return False + if padded_bs not in self.supported_batch_sizes: + return False + if padded_bs == num_requests: + return True # already at a captured size, nothing to add + # _get_padded_batch bails when padding_size + batch_size exceeds the + # configured batch size; that sum is just padded_bs. + return padded_bs <= self.config.batch_size + def _get_padded_batch(self, batch: ScheduledRequests, resource_manager: ResourceManager, runtime_draft_len: int) -> int: @@ -1017,11 +1207,19 @@ def _get_padded_batch(self, batch: ScheduledRequests, batch_size = batch.batch_size new_batch_size = batch_size - if self.enabled and self.config.enable_attention_dp and self.config.mapping.tp_size > 1: - # Padding only agrees on eligibility and batch size; the tier slot - # is unused here. It is still sent so both per-iteration exchanges - # share one fixed-size payload, and a constant keeps this path from - # resolving a tier for a batch it is about to change. + agreement = self.adp_shape_agreement + reuse_agreement = bool( + agreement is not None and agreement.reuse_graph_shape + and agreement.matches(batch, agreement.iteration, padded=False)) + if reuse_agreement: + if not can_run_cuda_graph or not agreement.padding_ready: + raise RuntimeError( + "DSpark cached ADP padding agreement lost a local " + "eligibility invariant") + new_batch_size = agreement.widest_batch_size + + if (self.enabled and self.config.enable_attention_dp + and self.config.mapping.tp_size > 1 and not reuse_agreement): graph_batch_info = self._gather_adp_graph_batch_info( batch, can_run_cuda_graph, SampleType.FULL) all_can_run_cuda_graph = all(rank_info[0] @@ -1041,6 +1239,11 @@ def _get_padded_batch(self, batch: ScheduledRequests, # otherwise this reduces to plain batch-size rounding. padded_batch_size = self._round_up_batch_size_with_draft_len( new_batch_size, runtime_draft_len) + if (reuse_agreement + and padded_batch_size != agreement.graph_batch_size): + raise RuntimeError( + "DSpark cached ADP padding agreement selected a different " + "graph batch size") if batch_size == padded_batch_size: return 0 @@ -1065,12 +1268,55 @@ def _get_padded_batch(self, batch: ScheduledRequests, key=f"cuda_graph_padding_dummy_fallback_{runtime_draft_len}") return 0 - batch.generation_requests.extend([padding_dummy_request] * padding_size) + # The ragged fit decides padding windows before this generic path + # appends cached dummy requests. A zero-real exact cell may use a low + # and a high object to represent a non-divisible V exactly. + secondary_padding_dummy = None + high_padding_rows = 0 + if self.ragged_zero_real_high_rows: + if (batch_size != 1 + or not batch.generation_requests[0].is_attention_dp_dummy + or self.ragged_zero_real_high_rows > padded_batch_size): + raise RuntimeError( + "DSpark zero-real heterogeneous padding lost its single " + "scheduled attention-DP dummy invariant") + high_padding_rows = self.ragged_zero_real_high_rows - 1 + if high_padding_rows > 0: + secondary_padding_dummy = self._get_or_create_padding_dummy( + resource_manager, runtime_draft_len, variant=1) + if secondary_padding_dummy is None: + logger.warning_once( + "Failed to allocate the secondary DSpark CUDA graph " + "padding dummy; falling back to native verification.", + key=("dspark_secondary_padding_dummy_fallback_" + f"{runtime_draft_len}")) + return 0 + if getattr(self.spec_config, "enable_confidence_scheduling", False): + ragged = any( + getattr(request, "py_verify_len", None) is not None + for request in batch.generation_requests) + padding_dummy_request.py_verify_len = (self.ragged_pad_verify_len + if ragged else None) + if secondary_padding_dummy is not None: + secondary_padding_dummy.py_verify_len = ( + self.ragged_pad_verify_len + 1 if ragged else None) + + low_padding_rows = padding_size - high_padding_rows + if low_padding_rows < 0: + raise RuntimeError( + "DSpark zero-real high-row count exceeds CUDA padding rows") + if high_padding_rows: + batch.generation_requests.extend([secondary_padding_dummy] * + high_padding_rows) + batch.generation_requests.extend([padding_dummy_request] * + low_padding_rows) return padding_size - def _get_or_create_padding_dummy( - self, resource_manager: ResourceManager, - runtime_draft_len: int) -> Optional[LlmRequest]: + def _get_or_create_padding_dummy(self, + resource_manager: ResourceManager, + runtime_draft_len: int, + *, + variant: int = 0) -> Optional[LlmRequest]: """Returns the padding dummy request for the given draft length, allocating it from the KV cache manager on first use. @@ -1078,8 +1324,13 @@ def _get_or_create_padding_dummy( (e.g. saturated cache); the batch cannot be padded until a retry succeeds. """ - if runtime_draft_len in self.padding_dummy_requests: - return self.padding_dummy_requests[runtime_draft_len] + if variant not in (0, 1): + raise ValueError( + f"unsupported CUDA padding dummy variant {variant}") + cache = (self.padding_dummy_requests + if variant == 0 else self.secondary_padding_dummy_requests) + if runtime_draft_len in cache: + return cache[runtime_draft_len] kv_cache_manager = resource_manager.get_resource_manager( self.config.kv_cache_manager_key) @@ -1104,8 +1355,10 @@ def _get_or_create_padding_dummy( draft_kv_cache_manager = get_draft_kv_cache_manager( self.spec_config, resource_manager) - # Use unique dummy request ID per draft length - dummy_request_id = CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len + max_draft_len = _padding_dummy_draft_token_capacity( + self.spec_config, runtime_draft_token_buffer_width) + dummy_request_id = cuda_graph_dummy_request_id( + runtime_draft_len, variant=variant, max_draft_len=max_draft_len) dummy_request = kv_cache_manager.add_dummy_requests( [dummy_request_id], token_nums=[ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM] @@ -1122,17 +1375,43 @@ def _get_or_create_padding_dummy( return None dummy_request = dummy_request[0] dummy_request.is_cuda_graph_dummy = True - if self.is_encoder_decoder: - if not self._add_cross_dummy_request( - dummy_request, resource_manager, dummy_encoder_output_len, - draft_kv_cache_manager): - return None + registered_managers: List[BaseResourceManager] = [kv_cache_manager] + if (draft_kv_cache_manager is not None + and draft_kv_cache_manager is not kv_cache_manager): + registered_managers.append(draft_kv_cache_manager) - spec_res_mgr = resource_manager.get_resource_manager( - ResourceManagerType.SPEC_RESOURCE_MANAGER) - if spec_res_mgr: - spec_res_mgr.add_dummy_requests([dummy_request_id]) - self.padding_dummy_requests[runtime_draft_len] = dummy_request + def rollback() -> None: + for manager in reversed(registered_managers): + manager.free_resources(dummy_request) + + try: + if self.is_encoder_decoder: + cross_kv_cache_manager = resource_manager.get_resource_manager( + ResourceManagerType.CROSS_KV_CACHE_MANAGER) + if (cross_kv_cache_manager is None + or not self._add_cross_dummy_request( + dummy_request, resource_manager, + dummy_encoder_output_len, draft_kv_cache_manager)): + rollback() + return None + if not any(cross_kv_cache_manager is manager + for manager in registered_managers): + registered_managers.append(cross_kv_cache_manager) + + spec_res_mgr = resource_manager.get_resource_manager( + ResourceManagerType.SPEC_RESOURCE_MANAGER) + if spec_res_mgr: + spec_res_mgr.add_dummy_requests([dummy_request_id]) + if not any(spec_res_mgr is manager + for manager in registered_managers): + registered_managers.append(spec_res_mgr) + except NoFreeSlotsError: + rollback() + return None + except Exception: + rollback() + raise + cache[runtime_draft_len] = dummy_request return dummy_request def _padding_dummy_managers( @@ -1169,20 +1448,33 @@ def release_padding_dummy(self, resource_manager: ResourceManager, manager that allocated part of it, and drops it from the runner so a later padded step re-creates it. - One dummy request ID is spread across up to four managers -- the main - KV cache manager, the one-model draft KV cache manager, the + Each low/high dummy request ID is spread across up to four managers -- + the main KV cache manager, the one-model draft KV cache manager, the speculative resource manager slot and the encoder-decoder cross-KV - cache manager. Releasing only the main one leaves the others holding - the ID, and re-creation reuses the same - ``CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len``. + cache manager. Releasing only the main one leaves the others holding + the ID; re-creation deterministically reuses the same variant ID. - Returns True if a dummy was held for that draft length. + Returns True if either member of the low/high dummy pair was held for + that draft length. """ - dummy_request = self.padding_dummy_requests.pop(runtime_draft_len, None) - if dummy_request is None: + dummy_requests = [ + cache.pop(runtime_draft_len, None) for cache in ( + self.padding_dummy_requests, + self.secondary_padding_dummy_requests, + ) + ] + dummy_requests = [ + request for request in dummy_requests if request is not None + ] + if not dummy_requests: return False - for manager in self._padding_dummy_managers(resource_manager): - manager.free_resources(dummy_request) + # The cached ADP agreement authenticated this dummy's presence. A + # rebalance may release it between iterations; invalidate the local + # agreement before another batch can suppress the padding collective. + self.adp_shape_agreement = None + for dummy_request in dummy_requests: + for manager in self._padding_dummy_managers(resource_manager): + manager.free_resources(dummy_request) return True def _can_pad_any_batch(self, runtime_draft_len: int) -> bool: @@ -1241,6 +1533,19 @@ def preallocate_padding_dummies(self, logger.info( "Pre-allocated the CUDA graph padding dummy request " f"(draft_len={draft_len}) at warmup.") + if (self._dspark_trims_submitted_tokens and getattr( + self.spec_config, "confidence_sps_table_path", None)): + if self._get_or_create_padding_dummy(resource_manager, + draft_len, + variant=1) is None: + logger.warning( + "Could not pre-allocate the secondary DSpark CUDA " + "graph padding dummy " + f"(draft_len={draft_len}); non-divisible zero-real " + "exact cells will remain native.") + else: + logger.info("Pre-allocated the secondary DSpark CUDA graph " + f"padding dummy (draft_len={draft_len}).") def _add_cross_dummy_request( self, dummy_request: LlmRequest, resource_manager: ResourceManager, @@ -1264,11 +1569,6 @@ def _add_cross_dummy_request( if cross_dummy_requests is not None: return True - kv_cache_manager = resource_manager.get_resource_manager( - self.config.kv_cache_manager_key) - kv_cache_manager.free_resources(dummy_request) - if draft_kv_cache_manager is not None: - draft_kv_cache_manager.free_resources(dummy_request) return False @staticmethod @@ -1339,6 +1639,9 @@ def clear(self): self.graph_metadata.clear() self._pre_replay_position_ids = None self.padding_dummy_requests = {} + self.secondary_padding_dummy_requests = {} + self.ragged_zero_real_high_rows = 0 + self.adp_shape_agreement = None del self.memory_pool self.memory_pool = None torch.cuda.empty_cache() diff --git a/tensorrt_llm/_torch/pyexecutor/dspark_mixed_host_capacity.py b/tensorrt_llm/_torch/pyexecutor/dspark_mixed_host_capacity.py new file mode 100644 index 000000000000..ab49919ba7c1 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/dspark_mixed_host_capacity.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pure host mixed geometry with context/full-width capacity reservations. + +The caller must authenticate the current ready host survival matrix and retain +request ownership. This module neither reads device tensors nor prices routes. +Actual allocation is delegated to the existing DSpark top-k implementation. +""" + +from dataclasses import dataclass +from typing import Callable, Tuple + +import numpy as np +import torch + + +@dataclass(frozen=True) +class Capacity: + proposed_body_rows: int + eligible_rows: int + optional_budget: int + total_real_rows: int + expected_yield: float + + +@dataclass(frozen=True) +class LocalCurve: + request_ids: Tuple[int, ...] + survival: torch.Tensor + context_rows: int + fixed_generation_rows: int + physical_k: int + min_verify_len: int + native_expected_yield: float + candidates: Tuple[Capacity, ...] + + def candidate(self, body_rows: int) -> Capacity: + for item in self.candidates: + if item.proposed_body_rows == body_rows: + return item + raise ValueError("Unlisted or locally infeasible body capacity") + + +def _nonnegative_integer(value: int, name: str) -> None: + if type(value) is not int or value < 0: + raise ValueError(f"{name} must be a nonnegative integer") + + +def prepare_local_curve( + *, + survival: torch.Tensor, + request_ids: Tuple[int, ...], + context_rows: int, + fixed_generation_rows: int, + physical_k: int, + min_verify_len: int, + body_capacities: Tuple[int, ...], +) -> LocalCurve: + """Reserve context/full-width rows before filling eligible draft prefixes. + + Infeasible capacities are omitted. Context work contributes to paid rows, + not fictitious generation progress. Final-context progress is preserved by + the caller as a separate, candidate-independent component. + """ + for name, value in ( + ("context_rows", context_rows), + ("fixed_generation_rows", fixed_generation_rows), + ("physical_k", physical_k), + ("min_verify_len", min_verify_len), + ): + _nonnegative_integer(value, name) + if not 1 <= min_verify_len <= physical_k: + raise ValueError("Invalid draft prefix bounds") + if ( + not isinstance(survival, torch.Tensor) + or survival.device.type != "cpu" + or survival.ndim != 2 + or tuple(survival.shape) != (len(request_ids), physical_k) + ): + raise ValueError("Requires a ready CPU [eligible requests, physical K] matrix") + if any(type(item) is not int or item < 0 for item in request_ids) or len( + set(request_ids) + ) != len(request_ids): + raise ValueError("Request identities must be unique nonnegative integers") + if ( + not body_capacities + or any(type(item) is not int or item <= 0 for item in body_capacities) + or tuple(sorted(set(body_capacities))) != body_capacities + ): + raise ValueError("Body capacities must be positive, strictly increasing") + # Own the admitted host matrix: later copy-buffer reuse cannot change the + # capacity curve between the group vote and its allocation. + matrix = survival.detach().to(dtype=torch.float32).clone() + values = matrix.numpy().astype(np.float64, copy=False) + if ( + not np.isfinite(values).all() + or (values < 0).any() + or (values > 1).any() + or (np.diff(values, axis=1) > 0).any() + ): + raise ValueError("Survival must be finite probabilities and prefix-monotone") + num_real = len(request_ids) + floor_rows = num_real * (min_verify_len + 1) + native_rows = num_real * (physical_k + 1) + base_yield = float(num_real) + float(values[:, :min_verify_len].sum(dtype=np.float64)) + ordered = np.sort(values[:, min_verify_len:physical_k].reshape(-1))[::-1] + prefix = np.concatenate(([0.0], np.cumsum(ordered, dtype=np.float64))) + native_yield = base_yield + float(prefix[-1]) + reserve = context_rows + fixed_generation_rows + candidates = [] + for body_rows in body_capacities: + available = body_rows - reserve + if available < floor_rows: + continue + eligible_rows = min(available, native_rows) + optional_budget = eligible_rows - floor_rows + candidates.append( + Capacity( + body_rows, + eligible_rows, + optional_budget, + reserve + eligible_rows, + base_yield + float(prefix[optional_budget]), + ) + ) + return LocalCurve( + tuple(request_ids), + matrix, + context_rows, + fixed_generation_rows, + physical_k, + min_verify_len, + native_yield, + tuple(candidates), + ) + + +def allocate_candidate( + curve: LocalCurve, + body_rows: int, + original_allocate: Callable[[torch.Tensor, int, int, int], torch.Tensor], +): + """Use original DSpark top-k with survival_eps=0 and validate exact rows.""" + candidate = curve.candidate(body_rows) + lengths = original_allocate( + curve.survival, candidate.optional_budget, curve.physical_k, curve.min_verify_len + ) + if ( + not isinstance(lengths, torch.Tensor) + or lengths.device.type != "cpu" + or lengths.ndim != 1 + or len(lengths) != len(curve.request_ids) + ): + raise ValueError("Allocator returned incompatible request lengths") + if lengths.dtype not in (torch.int32, torch.int64): + raise ValueError("Allocator lengths must be integer") + windows = tuple(int(value) for value in lengths.tolist()) + if ( + any(not curve.min_verify_len <= value <= curve.physical_k for value in windows) + or sum(value + 1 for value in windows) != candidate.eligible_rows + ): + raise ValueError("Allocator did not spend the admitted eligible-row budget") + return tuple(zip(curve.request_ids, windows)) + + +def realized_paid_bucket( + curves: Tuple[LocalCurve, ...], proposed_body_rows: int, captured_body_rows: Tuple[int, ...] +) -> int: + """Canonicalize actual all-rank maximum; do not manufacture padding work.""" + if not curves or proposed_body_rows not in captured_body_rows: + raise ValueError("Requires a captured proposal and nonempty rank group") + if ( + any(type(item) is not int or item <= 0 for item in captured_body_rows) + or tuple(sorted(set(captured_body_rows))) != captured_body_rows + ): + raise ValueError("Captured body grid must be positive and strictly increasing") + needed = max(curve.candidate(proposed_body_rows).total_real_rows for curve in curves) + for bucket in captured_body_rows: + if bucket >= needed: + return bucket + raise ValueError("Actual rows exceed the captured body grid") diff --git a/tensorrt_llm/_torch/pyexecutor/dspark_mixed_host_snapshot.py b/tensorrt_llm/_torch/pyexecutor/dspark_mixed_host_snapshot.py new file mode 100644 index 000000000000..4a86a012f92c --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/dspark_mixed_host_snapshot.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Current ready host snapshot ownership; decode device mode is untouched.""" + +from dataclasses import dataclass +from typing import Callable, Tuple + +import torch + +from .dspark_mixed_host_capacity import LocalCurve, prepare_local_curve + + +@dataclass(frozen=True) +class AuthenticatedCurve: + curve: LocalCurve + request_slots: Tuple[Tuple[int, int], ...] + producing_iteration: int + snapshot_staged_iteration: int + staging_sequence: int + ranking_mode: str = "mixed_current_ready_host" + request_incarnations: Tuple[int, ...] = () + + +def prepare_current_mixed_curve( + *, + planner, + request_ids: Tuple[int, ...], + rows: Tuple[int, ...], + iteration: int, + staging_sequence: int, + context_rows: int, + fixed_generation_rows: int, + body_capacities: Tuple[int, ...], + original_compute_survival: Callable, + requests: Tuple[object, ...] = (), + incarnations: Tuple[int, ...] = (), +) -> AuthenticatedCurve | None: + """Return an owned current-host curve, or native fallback without waiting. + + This is an explicit mixed host-window contract, not fresh device ranking. + It uses the existing ready/event and request-addressed gather API even when + decode planner.device_windows is true. It never toggles that global flag. + Source/runtime admission and all-rank voting belong to the caller. + """ + if ( + type(iteration) is not int + or iteration < 0 + or type(staging_sequence) is not int + or staging_sequence < 0 + or len(request_ids) != len(rows) + ): + return None + k, floor = int(planner.max_verify_len), int(planner.cfg.min_verify_len) + if not request_ids: + # A context/dummy-only rank has no invented confidence ownership/yield. + curve = prepare_local_curve( + survival=torch.empty((0, k), dtype=torch.float32), + request_ids=(), + context_rows=context_rows, + fixed_generation_rows=fixed_generation_rows, + physical_k=k, + min_verify_len=floor, + body_capacities=body_capacities, + ) + return AuthenticatedCurve(curve, (), iteration, -1, staging_sequence) + meta = getattr(planner, "_mixed_current_snapshot_meta", None) + snapshot = planner._ready_snapshot() # Existing query(), never synchronize(). + if ( + snapshot is None + or meta is None + or meta.get("buffer") is not snapshot + or meta.get("event") is not planner._copy_event + or meta.get("staging_sequence") != staging_sequence + or type(meta.get("staged_iteration")) is not int + or not 0 <= meta["staged_iteration"] < iteration + or snapshot.device.type != "cpu" + ): + return None + owners = meta.get("owners") + if ( + not isinstance(owners, dict) + or len(set(request_ids)) != len(request_ids) + or len(set(rows)) != len(rows) + or any( + type(request_id) is not int + or request_id < 0 + or type(row) is not int + or row < 0 + or row >= snapshot.shape[0] + or owners.get(request_id) != row + for request_id, row in zip(request_ids, rows) + ) + ): + return None + # Ownership consistency above remains fail-closed. Freshness is a separate + # conservative partition: no writer proof means full K, never fake scores. + stamp = meta.get("confidence_stamp") + attempts = meta.get("producer_attempts") + expected_writer = staging_sequence - 1 # Copy precedes forward q. + valid = [] + if ( + expected_writer > 0 + and len(requests) == len(rows) + and len(incarnations) == len(rows) + and isinstance(attempts, dict) + and stamp is getattr(planner, "_host_confidence_stamp", None) + and isinstance(stamp, torch.Tensor) + and stamp.device.type == "cpu" + and stamp.dtype == torch.int32 + and stamp.ndim == 1 + and stamp.shape[0] == snapshot.shape[0] + ): + for index, (request_id, row, request, incarnation) in enumerate( + zip(request_ids, rows, requests, incarnations) + ): + attempt = attempts.get(request_id) + if ( + type(incarnation) is int + and incarnation > 0 + and type(attempt) is tuple + and len(attempt) == 3 + and attempt[0] is request + and type(attempt[1]) is int + and type(attempt[2]) is int + and attempt[1] == row + and attempt[2] == incarnation + and int(stamp[row]) == expected_writer + ): + valid.append(index) + fixed_generation_rows += (len(rows) - len(valid)) * (k + 1) + request_ids = tuple(request_ids[index] for index in valid) + rows = tuple(rows[index] for index in valid) + incarnations = tuple(incarnations[index] for index in valid) + if not rows: + curve = prepare_local_curve( + survival=torch.empty((0, k), dtype=torch.float32), + request_ids=(), + context_rows=context_rows, + fixed_generation_rows=fixed_generation_rows, + physical_k=k, + min_verify_len=floor, + body_capacities=body_capacities, + ) + return AuthenticatedCurve(curve, (), iteration, meta["staged_iteration"], staging_sequence) + selected = planner._gather_rows(num_gen_requests=len(rows), rows=rows, snapshot=snapshot) + if selected is None: + return None + survival = original_compute_survival(planner.apply_calibration(selected)) + try: + curve = prepare_local_curve( + survival=survival, + request_ids=tuple(request_ids), + context_rows=context_rows, + fixed_generation_rows=fixed_generation_rows, + physical_k=k, + min_verify_len=floor, + body_capacities=body_capacities, + ) + except ValueError: + return None + return AuthenticatedCurve( + curve, + tuple(zip(request_ids, rows)), + iteration, + meta["staged_iteration"], + staging_sequence, + request_incarnations=incarnations, + ) diff --git a/tensorrt_llm/_torch/pyexecutor/dspark_mixed_runtime.py b/tensorrt_llm/_torch/pyexecutor/dspark_mixed_runtime.py new file mode 100644 index 000000000000..3873b8adf2d7 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/dspark_mixed_runtime.py @@ -0,0 +1,474 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Admitted host-mixed allocation in the existing DSpark policy transaction. + +There is no prescribed action, diagnostic pricing mode, new collective, or +decode-table conversion. A missing quote/owner/route keeps mixed steps native. +""" + +import math +from dataclasses import dataclass, replace +from fractions import Fraction + +from ..speculative.dspark_mixed_evidence import AdmittedMixedCosts, MixedGeometry +from ..speculative.dspark_planner import _canonical_json_sha256 +from ..speculative.dspark_schedule import ( + DSparkScheduleConfig, + compute_survival, + schedule_verify_lens_topk, +) +from .dspark_mixed_host_capacity import allocate_candidate, prepare_local_curve +from .dspark_mixed_host_snapshot import AuthenticatedCurve, prepare_current_mixed_curve + +BODY_BUCKETS = (96, 128, 192, 256, 384, 512) +WIRE_WORDS = 46 +_MAGIC = 0x4D485031 +_VERSION = 1 +_K = 3 +_WORLD = 8 +_SCALE = 1_000_000 + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise ValueError(message) + + +def _integer(value: int) -> int: + _require(type(value) is int and 0 <= value < 2**63, "Invalid mixed geometry integer") + return value + + +def _membership(batch) -> tuple: + def rows(requests): + return tuple((id(r), r.py_request_id, r.py_seq_slot, bool(r.is_dummy)) for r in requests) + + return rows(batch.context_requests), rows(batch.generation_requests) + + +def _geometry(batch) -> tuple: + contexts = tuple( + ( + id(r), + r.context_chunk_size, + r.context_current_position, + r.py_num_compressed_tokens, + r.py_max_new_tokens, + bool(r.is_first_context_chunk), + ) + for r in batch.context_requests + ) + generations = tuple( + ( + id(r), + r.py_batch_idx, + len(r.py_draft_tokens), + r.get_num_tokens(0), + getattr(r.state, "name", None), + r.py_orig_prompt_len, + r.py_max_new_tokens, + ) + for r in batch.generation_requests + ) + return contexts, generations, tuple(id(r) for r in batch.context_requests_last_chunk) + + +def _remaining(request, runner) -> int: + return max( + 0, + min( + _integer(int(request.py_orig_prompt_len)) + _integer(int(request.py_max_new_tokens)), + runner._config.max_seq_len, + ) + - _integer(int(request.get_num_tokens(0))), + ) + + +def _greedy(batch) -> bool: + from tensorrt_llm.sampling_params import SamplingParams + + from .sampler.sampler_common import _request_get_sampling_params + + return all( + SamplingParams.params_imply_greedy_decoding( + temperature=p.temperature, + top_p=p.top_p, + top_k=p.top_k, + min_p=p.min_p, + top_p_decay=p.top_p_decay, + use_beam_search=p.use_beam_search, + ) + for p in (_request_get_sampling_params(r) for r in batch.all_requests()) + ) + + +@dataclass(frozen=True) +class _Ticket: + batch: object + membership: tuple + geometry: tuple + iteration: int + rank: int + eligible: tuple + curve: AuthenticatedCurve + vote: tuple[int, ...] + + +def stage_mixed_ownership(planner, worker, iteration: int, generation_requests) -> None: + """Bind the existing asynchronous copy, without reading device storage. + + An attempted forward is not a write. At the next decision, the copied + positive producer stamp must match that attempt and the current incarnation. + """ + if not planner._snapshot_valid or planner._host_buffer is None: + planner._mixed_current_snapshot_meta = None + return + previous = getattr(planner, "_mixed_current_snapshot_meta", None) + sequence = int(worker._draft_seq_host) + attempts = {} + for request in generation_requests: + if request.is_dummy: + continue + request_id = int(request.py_request_id) + slot = worker._req_to_slot.get(request_id) + incarnation = worker.confidence_incarnation_for(request_id) + if type(slot) is int and type(incarnation) is int and incarnation > 0: + attempts[request_id] = (request, slot, incarnation) + producer_attempts = ( + previous.get("generation_attempts", {}) + if isinstance(previous, dict) and previous.get("staging_sequence") == sequence - 1 + else {} + ) + planner._mixed_current_snapshot_meta = dict( + buffer=planner._host_buffer, + event=planner._copy_event, + owners=dict(worker._req_to_slot), + staged_iteration=int(iteration), + staging_sequence=sequence, + confidence_stamp=planner._host_confidence_stamp, + generation_attempts=attempts, + producer_attempts=producer_attempts, + ) + + +class MixedRuntime: + """A separately bound optional mixed path; decode admission is untouched.""" + + def __init__(self, runner, costs: AdmittedMixedCosts | None) -> None: + self.runner = runner + self.costs = costs + self._ticket: _Ticket | None = None + identity = ( + None + if costs is None + else { + "live": costs.identity.payload(), + "costs": costs.binding.costs_sha256, + "receipt": costs.binding.admission_sha256, + "protocol": costs.binding.protocol_sha256, + } + ) + digest = "0" * 64 if identity is None else _canonical_json_sha256(identity) + self._identity_words = tuple(int(digest[n : n + 8], 16) for n in range(0, 64, 8)) + + def reset_step(self) -> None: + self._ticket = None + self.runner._dspark_host_window_step = False + self.runner.cuda_graph_runner._dspark_host_window_batch = None + + def _route_ready(self, batch, planner) -> bool: + cfg, mapping = self.runner._config, self.runner.mapping + spec, breakable = cfg.spec_config, self.runner.breakable_cuda_graph_runner + return bool( + self.costs is not None + and self.runner._dspark_confidence_enabled + and self.runner._dspark_trims_submitted_tokens + and cfg.max_draft_len == spec.max_draft_len == spec.block_size == _K + and spec.decoding_type == "DSpark" + and spec.draft_is_embedded_in_target + and not getattr(spec, "use_rejection_sampling", False) + and mapping.tp_size == _WORLD + and mapping.pp_size == mapping.cp_size == 1 + and cfg.enable_attention_dp + and not cfg.torch_compile_enabled + and not self.runner.use_beam_search + and not cfg.is_multimodal + and not getattr(batch, "encoder_requests", ()) + and getattr(cfg.prefill_cuda_graph_backend, "name", None) == "BREAKABLE" + and tuple(sorted(cfg.prefill_cuda_graph_num_tokens)) == BODY_BUCKETS + and breakable is not None + and not breakable.is_capturing + and not breakable.is_warming_up + and all(breakable.has_graph(bucket) for bucket in BODY_BUCKETS) + and planner.max_verify_len == _K + and planner.cfg.min_verify_len == 1 + and _greedy(batch) + ) + + def local_vote( + self, batch, worker, planner, iteration: int, rank: int, *, executor_eligible: bool + ) -> list[int]: + """Append a fixed suffix to the existing allgather; never wait for a copy.""" + out = [ + _MAGIC, + _VERSION, + int(rank), + int(iteration), + sum(not r.is_dummy for r in batch.context_requests), + 0, + *self._identity_words, + *([0] * 32), + ] + _require(len(out) == WIRE_WORDS, "Mixed vote extent differs") + self._ticket = None + try: + _require( + executor_eligible and self._route_ready(batch, planner), "Unsupported mixed route" + ) + contexts, generations = list(batch.context_requests), list(batch.generation_requests) + # First-chunk KV preparation can resolve prefix-cache hits and + # change context position/chunk size after this policy transaction. + _require( + not any(r.is_first_context_chunk for r in contexts), + "Mixed context geometry is not finalized", + ) + native_rows = sum(_integer(r.context_chunk_size) for r in contexts) + native_rows += len(generations) * (_K + 1) + _require( + native_rows <= _integer(self.runner._config.max_num_tokens), + "Mixed native rows exceed the prepared token budget", + ) + real = [r for r in generations if not r.is_dummy] + _require( + all( + getattr(r.state, "name", None) not in (None, "GENERATION_COMPLETE") + and _remaining(r, self.runner) > 0 + for r in real + ), + "Completed generation", + ) + meta = getattr(planner, "_mixed_current_snapshot_meta", None) + owners = meta.get("owners") if isinstance(meta, dict) else None + prospective = [ + r for r in real if r.py_batch_idx is not None and len(r.py_draft_tokens) >= _K + ] + _require( + len({r.py_request_id for r in prospective}) == len(prospective), "Duplicate owner" + ) + eligible = [ + r for r in prospective if not isinstance(owners, dict) or r.py_request_id in owners + ] + fixed = [r for r in generations if id(r) not in {id(e) for e in eligible}] + chunks = tuple(_integer(r.context_chunk_size) for r in contexts) + _require(all(chunks), "Empty context chunk") + curve = prepare_current_mixed_curve( + planner=planner, + request_ids=tuple(int(r.py_request_id) for r in eligible), + rows=tuple(worker.confidence_row_for(r.py_request_id) for r in eligible), + iteration=iteration, + staging_sequence=int(worker._draft_seq_host), + context_rows=sum(chunks), + fixed_generation_rows=len(fixed) * (_K + 1), + body_capacities=BODY_BUCKETS, + original_compute_survival=compute_survival, + requests=tuple(eligible), + incarnations=tuple( + worker.confidence_incarnation_for(r.py_request_id) for r in eligible + ), + ) + _require(curve is not None, "Current owned snapshot unavailable") + eligible = [r for r in eligible if r.py_request_id in curve.curve.request_ids] + fixed = [r for r in generations if id(r) not in {id(e) for e in eligible}] + useful = curve.curve.survival.clone() + for row, request in enumerate(eligible): + useful[row, max(0, _remaining(request, self.runner) - 1) :] = 0 + curve = replace( + curve, + curve=prepare_local_curve( + survival=useful, + request_ids=curve.curve.request_ids, + context_rows=sum(chunks), + fixed_generation_rows=len(fixed) * (_K + 1), + physical_k=_K, + min_verify_len=1, + body_capacities=BODY_BUCKETS, + ), + ) + real_fixed = [r for r in fixed if not r.is_dummy] + last = {id(r) for r in batch.context_requests_last_chunk if not r.is_dummy} + _require( + last <= {id(r) for r in contexts if not r.is_dummy}, + "Final-context ownership differs", + ) + context_progress = sum(r.py_max_new_tokens > 0 for r in contexts if id(r) in last) + lower = context_progress + len(real_fixed) + upper = context_progress + sum( + min(_K + 1, _remaining(r, self.runner)) for r in real_fixed + ) + histories = [_integer(int(r.get_num_tokens(0))) for r in real] + out[14:26] = [ + sum(chunks), + sum(n * n for n in chunks), + sum(n * _integer(r.context_current_position) for n, r in zip(chunks, contexts)), + sum(n * _integer(r.py_num_compressed_tokens) for n, r in zip(chunks, contexts)), + sum(histories), + max(histories, default=0), + len(real_fixed), + out[4], + len(eligible), + len(contexts), + len(generations), + len(fixed) * (_K + 1), + ] + out[26:28] = [ + math.ceil((curve.curve.native_expected_yield + upper) * _SCALE), + native_rows, + ] + for index, bucket in enumerate(BODY_BUCKETS): + try: + cell = curve.curve.candidate(bucket) + except ValueError: + continue + out[28 + 3 * index : 31 + 3 * index] = [ + 1, + math.floor((cell.expected_yield + lower) * _SCALE), + cell.total_real_rows, + ] + out[5] = 1 + _require(all(type(n) is int and 0 <= n < 2**63 for n in out), "Vote integer overflow") + self._ticket = _Ticket( + batch, + _membership(batch), + _geometry(batch), + iteration, + rank, + tuple(eligible), + curve, + tuple(out), + ) + except ( + ValueError, + TypeError, + KeyError, + AttributeError, + RuntimeError, + IndexError, + OverflowError, + ): + out[5] = 0 + self._ticket = None + # Invalid optional geometry cannot poison the collective's int64 + # extent. Its real-context count still selects group-native mixed. + out[14:] = [0] * (WIRE_WORDS - 14) + return out + + def split_votes(self, payloads, legacy_words: int) -> tuple[list, list]: + _require( + len(payloads) == _WORLD and all(len(p) == legacy_words + WIRE_WORDS for p in payloads), + "Mixed policy suffix shape differs", + ) + peers = [list(p[legacy_words:]) for p in payloads] + for rank, row in enumerate(peers): + _require( + all(type(n) is int and 0 <= n < 2**63 for n in row) + and row[:3] == [_MAGIC, _VERSION, rank] + and row[5] in (0, 1), + "Mixed policy header differs", + ) + for index, bucket in enumerate(BODY_BUCKETS): + feasible, amount, total = row[28 + 3 * index : 31 + 3 * index] + _require( + feasible in (0, 1) + and (feasible or amount == total == 0) + and (not feasible or total <= bucket and amount <= row[26]), + "Mixed capacity vote differs", + ) + _require(len({r[3] for r in peers}) == 1, "Mixed iteration differs") + return [list(p[:legacy_words]) for p in payloads], peers + + def apply(self, batch, worker, planner, peers, iteration: int) -> bool: + """Return False only for all-decode; preserve its original exact policy.""" + if not sum(row[4] for row in peers): + return False + _require(all(row[3] == iteration for row in peers), "Mixed publication iteration differs") + chosen, windows = None, {} + if self.costs is not None and all( + row[5] and tuple(row[6:14]) == self._identity_words for row in peers + ): + features = tuple(tuple(row[14:26]) for row in peers) + native_rows = tuple(row[27] for row in peers) + native_yield = sum(row[26] for row in peers) + for index, capacity in enumerate(BODY_BUCKETS): + cells = [row[28 + 3 * index : 31 + 3 * index] for row in peers] + if not all(cell[0] for cell in cells): + continue + try: + geometry = MixedGeometry( + features, capacity, tuple(c[2] for c in cells), native_rows + ) + except ValueError: + continue + quote = self.costs.lookup(geometry) + if quote is None: + continue + candidate = Fraction(sum(cell[1] for cell in cells), 1) / Fraction( + str(quote.candidate_upper_ms) + ) + native = Fraction(native_yield, 1) / Fraction(str(quote.native_lower_ms)) + margin = 1 + Fraction(str(self.costs.minimum_predicted_gain)) + if candidate > native and candidate >= native * margin: + item = (candidate, -capacity, geometry) + if chosen is None or item[:2] > chosen[:2]: + chosen = item + if chosen is not None: + ticket = self._ticket + _require( + ticket is not None + and ticket.batch is batch + and ticket.iteration == iteration + and ticket.membership == _membership(batch) + and ticket.geometry == _geometry(batch) + and ticket.vote == tuple(peers[ticket.rank]), + "Mixed request/geometry ticket changed", + ) + _require( + worker._draft_seq_host == ticket.curve.staging_sequence + and tuple(worker.confidence_row_for(r.py_request_id) for r in ticket.eligible) + == tuple(row for _, row in ticket.curve.request_slots) + and tuple( + worker.confidence_incarnation_for(r.py_request_id) for r in ticket.eligible + ) + == ticket.curve.request_incarnations, + "Mixed snapshot/slot incarnation changed", + ) + + def original_allocate(survival, budget, k, floor): + return schedule_verify_lens_topk( + survival=survival, + budget=budget, + cfg=DSparkScheduleConfig(k, floor, survival_eps=0.0), + ) + + windows = dict( + allocate_candidate(ticket.curve.curve, chosen[2].capacity, original_allocate) + ) + actual = ticket.curve.curve.context_rows + ticket.curve.curve.fixed_generation_rows + actual += sum(value + 1 for value in windows.values()) + _require( + actual == chosen[2].actual_rank_rows[ticket.rank], "Mixed allocated rows differ" + ) + values = [ + _K if r.is_dummy else windows.get(r.py_request_id, _K) + for r in batch.generation_requests + ] + _require(all(type(v) is int and 1 <= v <= _K for v in values), "Mixed window bounds differ") + for request, value in zip(batch.generation_requests, values): + request.py_verify_len = value + graph = self.runner.cuda_graph_runner + graph.agreed_ragged_bucket = None + graph.ragged_pad_verify_len = 0 + graph.ragged_zero_real_high_rows = 0 + graph._dspark_host_window_batch = batch + self.runner._dspark_device_budget = None + self.runner._dspark_host_window_step = True + return True diff --git a/tensorrt_llm/_torch/pyexecutor/engine/metadata.py b/tensorrt_llm/_torch/pyexecutor/engine/metadata.py index c9ed43c82ef9..48068a897eb0 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/metadata.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/metadata.py @@ -53,6 +53,7 @@ def build_attention_metadata( draft_kv_cache_manager: KVCacheManager | KVCacheManagerV2 | None = None, enable_context_mla_with_cached_kv: bool | None = None, num_heads_per_kv: int | None = None, + enable_ragged_verification: bool = False, ) -> AttentionMetadata: """Construct attention metadata and resolve model-derived inputs.""" pretrained_config = model_config.pretrained_config @@ -69,6 +70,9 @@ def build_attention_metadata( if sparse_attention_config is not None else None ) + ragged_metadata_kwargs = {} + if hasattr(attention_backend.Metadata, "enable_ragged_verification"): + ragged_metadata_kwargs["enable_ragged_verification"] = enable_ragged_verification return attention_backend.Metadata( max_num_requests=max_batch_size, max_num_tokens=max_num_tokens, @@ -82,4 +86,5 @@ def build_attention_metadata( cache_indirection=cache_indirection, num_heads_per_kv=num_heads_per_kv, sparse_metadata_params=sparse_metadata_params, + **ragged_metadata_kwargs, ) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index f09883c04bcf..62abdfb38a2d 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -11,7 +11,7 @@ import os from collections.abc import Callable, Sequence from contextlib import contextmanager -from typing import Any +from typing import Any, List, Optional, Tuple import torch @@ -46,7 +46,11 @@ BaseMambaCacheManager, MambaHybridCacheManager, ) -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, get_draft_token_length +from tensorrt_llm._torch.pyexecutor.llm_request import ( + LlmRequest, + get_draft_token_length, + get_request_tokens_per_gen_step, +) from tensorrt_llm._torch.pyexecutor.resource_manager import ( BaseResourceManager, KVCacheManager, @@ -70,6 +74,7 @@ ) from tensorrt_llm._torch.speculative.dspark import DSparkSpecMetadata from tensorrt_llm._torch.speculative.interface import INVALID_PROMPT_LOOKAHEAD_TOKEN +from tensorrt_llm._torch.speculative.ragged_helpers import ragged_gather_index_lists from tensorrt_llm._torch.speculative.spec_sampler_base import SampleStateTensorsSpec from tensorrt_llm._torch.speculative.utils import get_static_draft_len, resolve_draft_len from tensorrt_llm._torch.utils import ( @@ -404,6 +409,14 @@ def _allocate_decoder_buffers(self) -> None: self.previous_batch_indices_cuda = torch.empty( (self._config.max_num_tokens,), dtype=torch.int, device="cuda" ) + self.ragged_verify_lens_cuda = torch.empty( + (self._config.max_batch_size + 1,), dtype=torch.int, device="cuda" + ) + self.ragged_qo_indptr_cuda = torch.empty( + (self._config.max_batch_size + 2,), dtype=torch.int, device="cuda" + ) + self._pinned_host_cache = {} + self._pinned_host_active = {} self.input_ids_cuda = torch.empty( (self._config.max_num_tokens,), dtype=torch.int, device="cuda" ) @@ -431,6 +444,79 @@ def _allocate_decoder_buffers(self) -> None: def _init_decoder_state(self) -> None: """Initialize the mutable state read only by engine decoder execution.""" + spec_config = self._config.spec_config + self._dspark_confidence_enabled = bool( + spec_config is not None and getattr(spec_config, "enable_confidence_scheduling", False) + ) + self._dspark_fused_scheduler_enabled = bool( + spec_config is not None + and getattr(spec_config, "enable_fused_confidence_scheduler", False) + ) + self._dspark_trims_submitted_tokens = self._dspark_confidence_enabled + self._dspark_device_windows = bool( + self._dspark_trims_submitted_tokens + and getattr(spec_config, "enable_fused_confidence_scheduler", False) + ) + self._dspark_device_window_workspace = None + self._dspark_host_window_step = False + self._dspark_mixed_runtime = None + mixed_config = getattr(spec_config, "confidence_mixed_evidence", None) + if ( + self._dspark_confidence_enabled + and mixed_config is not None + and self._config.max_draft_len == 3 + and self.mapping.tp_size == 8 + and self.mapping.pp_size == self.mapping.cp_size == 1 + and self._config.enable_attention_dp + ): + from tensorrt_llm._torch.pyexecutor.dspark_mixed_runtime import MixedRuntime + from tensorrt_llm._torch.speculative.dspark_mixed_evidence import load_mixed_evidence + + resolution = load_mixed_evidence(mixed_config) + self._dspark_mixed_runtime = MixedRuntime(self, resolution.costs) + logger.info(f"DSpark mixed evidence: {resolution.reason}") + if self._dspark_device_windows: + from tensorrt_llm._torch.speculative.dspark_device_select import DeviceWindowWorkspace + + self._dspark_device_window_workspace = DeviceWindowWorkspace.allocate( + max_rows=self._config.max_batch_size + 1, + max_tokens=self._config.max_num_tokens, + device=self.input_ids_cuda.device, + ) + self._dspark_sps_cost_table = None + self._dspark_sps_artifact_metadata = None + self._dspark_sps_observed_runtime = None + self._dspark_exact_candidate_cells = () + self._dspark_exact_identity_words = (0,) * 8 + table_path = getattr(spec_config, "confidence_sps_table_path", None) + if self._dspark_confidence_enabled and table_path: + from tensorrt_llm._torch.speculative.dspark_planner import load_runtime_sps_cost_table + + self._dspark_sps_cost_table, self._dspark_sps_artifact_metadata = ( + load_runtime_sps_cost_table( + table_path, + graph_batch_sizes=self._config.cuda_graph_batch_sizes, + max_draft_len=int(spec_config.max_draft_len), + live_engine_fingerprint_path=getattr( + spec_config, "confidence_sps_live_fingerprint_path", None + ), + validate_freshness=spec_config.validate_confidence_artifact_freshness, + ) + ) + self._dspark_sps_observed_runtime = { + "max_draft_len": int(spec_config.max_draft_len), + "rank_local_graph_batch_sizes": list(self._config.cuda_graph_batch_sizes), + "gpu": torch.cuda.get_device_name(self.input_ids_cuda.device), + "torch_version": torch.__version__, + "validate_artifact_freshness": spec_config.validate_confidence_artifact_freshness, + } + self._dspark_exact_candidate_cells = self._dspark_sps_cost_table.candidate_cells() + self._dspark_exact_identity_words = ( + self._dspark_sps_cost_table.collective_identity_words + ) + self.previous_request_ids = [] + self.previous_verify_lens = [] if self._dspark_trims_submitted_tokens else None + self.has_previous_device_draft = False self._eager_workspace_reclaimer: EagerWorkspaceReclaimer | None = None # Steady-state generation-only prepare cache (non-speculative overlap # decode). Holds the per-request lists that are invariant while the @@ -691,6 +777,8 @@ def _warmup_scheduled(self, resource_manager: ResourceManager, kv_cache_manager) # The lifetime of model engine and kv cache manager can be different. # Reset the global cuda graph dummy requests in warmup. self.cuda_graph_runner.padding_dummy_requests = {} + self.cuda_graph_runner.secondary_padding_dummy_requests = {} + self.cuda_graph_runner.ragged_zero_real_high_rows = 0 if self.mapping.cp_size > 1: cp_type = self.mapping.cp_config.get("cp_type", None) @@ -791,6 +879,8 @@ def _warmup_scheduled(self, resource_manager: ResourceManager, kv_cache_manager) finally: self.cuda_graph_runner.is_warmup_only = False self.cuda_graph_runner.padding_dummy_requests = {} + self.cuda_graph_runner.secondary_padding_dummy_requests = {} + self.cuda_graph_runner.ragged_zero_real_high_rows = 0 self._run_cuda_graph_warmup(resource_manager) log_mem_snapshot("warmup/after_cuda_graph_capture") # Pre-compile DeepGEMM paged_mqa_logits_metadata for every 32-aligned @@ -806,6 +896,8 @@ def _warmup_scheduled(self, resource_manager: ResourceManager, kv_cache_manager) # creates; build it when every forward above was skipped. with self._warmup_timer.phase("dsa_prewarm"): self._ensure_dsa_attn_metadata_for_warmup(resource_manager) + self._warmup_dspark_ragged_compressor_metadata() + self._warmup_dspark_fused_scheduler() with timing_metric("dg_paged_mqa_warmup_seconds", self._metrics): self._warmup_dg_paged_mqa_logits_metadata() log_mem_snapshot("warmup/after_dg_paged_mqa_logits_metadata") @@ -1853,7 +1945,9 @@ def _compute_dynamic_draft_len_mapping(self) -> dict[int, int] | None: mapping[graph_bs] = draft_len return mapping - def _get_graphs_to_capture(self, cuda_graph_batch_sizes: list[int]) -> list[tuple[int, int]]: + def _get_graphs_to_capture( + self, cuda_graph_batch_sizes: list[int] + ) -> list[tuple[int, int] | tuple[int, int, int]]: """Determine which (batch_size, draft_len) graphs to capture. Returns: @@ -1893,6 +1987,33 @@ def _get_graphs_to_capture(self, cuda_graph_batch_sizes: list[int]) -> list[tupl ) return graphs + if self._dspark_confidence_enabled: + max_draft_len = int(self._config.spec_config.max_draft_len) + native_graphs = [(bs, max_draft_len) for bs in cuda_graph_batch_sizes] + if not self._dspark_trims_submitted_tokens: + return native_graphs + from tensorrt_llm._torch.speculative.dspark_planner import ExactSpsCostTable + + exact_table = self._dspark_sps_cost_table + if not isinstance(exact_table, ExactSpsCostTable): + raise RuntimeError( + "DSpark confidence graph capture requires an authenticated exact SPS table" + ) + ragged_graphs = [ + (bs, max_draft_len, verifier_budget) + for bs in cuda_graph_batch_sizes + for verifier_budget in exact_table.production_candidate_budgets(bs) + ] + # Full K keeps the ordinary native graph. Compact layouts have + # distinct keys; they cannot alias its uniform attention layout. + graphs = native_graphs + ragged_graphs + logger.info( + f"DSpark ragged verification: capturing {len(graphs)} graphs " + f"({len(native_graphs)} native static-K{max_draft_len} graphs + " + f"{len(ragged_graphs)} exact measured (G,V) cells)." + ) + return graphs + # Case 2: Static draft length # Match the runtime_draft_len semantics enforced in _prepare_tp_inputs: # logical K for linear-tree modes, total tree tokens for tree decoding. @@ -2038,7 +2159,9 @@ def _run_capture_pass( pinned_tier = sample_type or SampleType.FULL self.cuda_graph_runner.set_capture_sample_type(pinned_tier) try: - for bs, draft_len in graphs_to_capture: + for entry in graphs_to_capture: + bs, draft_len = entry[:2] + verify_bucket = entry[2] if len(entry) == 3 else None if bs > self._config.max_batch_size: continue @@ -2051,6 +2174,8 @@ def _run_capture_pass( force_non_greedy=force_non_greedy, ) with self._release_batch_context(warmup_request, resource_manager) as batch: + if batch is not None and verify_bucket is not None: + self._set_warmup_ragged_windows(batch, verify_bucket, draft_len) if batch is None: # No KV cache space for this batch size. During KV # cache estimation this makes the profiling peak @@ -2087,6 +2212,7 @@ def _run_capture_pass( finally: self._force_lora_graph_for_capture = None self.cuda_graph_runner.set_capture_sample_type(None) + self.cuda_graph_runner.agreed_ragged_bucket = None if self._lora.cuda_graph_manager is None: lora_graph_cases = [False] @@ -2137,7 +2263,8 @@ def _capture_variant( # Dummy warmup requests carry no sampling params, so the greedy pass # needs no override while the advanced one has to force the flag. _capture_variant("greedy") - _capture_variant("advanced sampling", force_non_greedy=True) + if self._supports_advanced_draft_sampling(): + _capture_variant("advanced sampling", force_non_greedy=True) else: # TorchSampler branch: FULL carries no sampling in the graph and is # what every batch replays unless in-graph sampling is enabled, so it @@ -2691,6 +2818,7 @@ def _set_up_attn_metadata( else None, kv_cache_manager=kv_cache_manager, draft_kv_cache_manager=draft_kv_cache_manager, + enable_ragged_verification=self._dspark_confidence_enabled, ) if isinstance(kv_cache_manager, BaseMambaCacheManager): self.attn_metadata.mamba_chunk_size = getattr( @@ -2924,7 +3052,9 @@ def _sync_group_all_greedy_sample(self, spec_metadata) -> None: Why the sampling-path choice must be group-uniform under ADP + LM-head TP is documented on the anchor, - ``SpecMetadata.group_all_greedy_sample``. Local contract: called once + ``SpecMetadata.group_all_greedy_sample``. Confidence also needs a + uniform flag to reject unsupported advanced sampling on every rank. + Local contract: called once per iteration, right after ``update_is_all_greedy_sample`` and BEFORE the CUDA graph key is built. The gate is pure config (identical on every rank), so ranks also agree on whether the exchange happens; the @@ -2933,19 +3063,40 @@ def _sync_group_all_greedy_sample(self, spec_metadata) -> None: ``all_rank_num_tokens`` exchange, which runs in ``_prepare_inputs`` -- after the graph key, too late for the key to see the synced value. """ - # enable_lm_head_tp_in_adp implies enable_attention_dp (asserted in - # Mapping.__init__), so ADP needs no separate check here. - if not (self.mapping.enable_lm_head_tp_in_adp and spec_metadata.use_rejection_sampling): + # Confidence's unsupported advanced path must also be rejected by all + # ranks before graph replay, even without LM-head TP in attention DP. + # Feature-off engines retain the existing LM-head-TP gate. + if not ( + spec_metadata.use_rejection_sampling + and (self.mapping.enable_lm_head_tp_in_adp or self._dspark_confidence_enabled) + ): return local_flag = bool(spec_metadata.is_all_greedy_sample) - all_flags = self.dist.tp_allgather_int64([local_flag])[:, 0] - spec_metadata.group_all_greedy_sample = bool(all_flags.all()) + if self.mapping.tp_size > 1: + all_flags = self.dist.tp_allgather_int64([local_flag])[:, 0] + spec_metadata.group_all_greedy_sample = bool(all_flags.all()) + else: + spec_metadata.group_all_greedy_sample = local_flag # Also overwrite the live flag directly: this iteration's scan already # ran (update_is_all_greedy_sample just returned) and the CUDA graph # key reads the flag next -- the stored override only takes effect on # the NEXT rescan (populate), which is after key selection. spec_metadata.is_all_greedy_sample = spec_metadata.group_all_greedy_sample + def _supports_advanced_draft_sampling(self) -> bool: + """Whether the engine may capture an advanced draft-sampling variant.""" + return not ( + self._dspark_confidence_enabled and self._config.spec_config.use_rejection_sampling + ) + + def _validate_dspark_draft_sampling(self, spec_metadata: SpecMetadata) -> None: + """Reject an unsupported group-uniform sampler path before graph lookup.""" + if self._dspark_confidence_enabled and spec_metadata.wants_advanced_draft_sampling: + raise ValueError( + "DSpark confidence scheduling does not support advanced rejection draft sampling: " + "confidence is conditioned on the greedy proposal prefix" + ) + def _is_final_multimodal_context_decode_compatible(self, request: LlmRequest) -> bool: """Return whether the final prompt token uses the decode input path. @@ -3006,6 +3157,823 @@ def _record_cached_kv_tokens_per_req( num_cached_tokens_per_seq ) - sum(counts) + @torch.inference_mode() + def _warmup_dspark_fused_scheduler(self) -> None: + """Compile the optional policy-neutral scheduler before live requests. + + The device-window prologue is deliberately skipped during whole-model + CUDA-graph capture, so its Triton rank/finalize kernels would otherwise + compile on the first live compact step. Compile every authenticated + exact-SPS graph row count on every rank and record only successful + shapes. Unsupported shapes retain the established tensor fallback. + """ + self._dspark_fused_schedule_ready_sizes = set() + if ( + not self._dspark_fused_scheduler_enabled + or not self._dspark_confidence_enabled + or not self._dspark_trims_submitted_tokens + or not self._dspark_device_windows + ): + return + + worker = self._get_spec_worker() + planner = getattr(worker, "verify_planner", None) + exact_table = planner.exact_cost_table if planner is not None else None + if exact_table is None: + return + + from tensorrt_llm._torch.speculative.dspark_schedule import ( + schedule_verify_lens_topk_fused_fill, + ) + + cfg = planner.cfg + if cfg.survival_eps <= 0.0: + logger.warning_once( + "DSpark fused scheduler requires survival_eps > 0; using the " + "tensor scheduler for every graph size", + key="dspark_fused_scheduler_zero_epsilon", + ) + return + ready = [] + with torch.inference_mode(): + for graph_bs in sorted(int(g) for g in exact_table.tables): + if graph_bs > 256: + logger.warning_once( + "DSpark fused scheduler supports at most " + f"256 rows; G={graph_bs} will use the tensor fallback", + key=(f"dspark_fused_scheduler_unsupported_{graph_bs}"), + ) + continue + survival = torch.ones( + (graph_bs, int(cfg.block_size)), + dtype=torch.float32, + device="cuda", + ) + try: + warmup_budget = min(1, graph_bs * int(cfg.schedulable_per_request)) + graph_num_tokens = graph_bs * (int(cfg.min_verify_len) + 1) + warmup_budget + schedule_verify_lens_topk_fused_fill( + survival=survival, + budget=warmup_budget, + num_real=graph_bs, + pad_len=1, + cfg=cfg, + graph_num_tokens=graph_num_tokens, + ) + except Exception as exc: + logger.warning_once( + "DSpark fused scheduler prewarm failed for " + f"G={graph_bs}; using the tensor fallback. " + f"{type(exc).__name__}: {exc}", + key=(f"dspark_fused_scheduler_prewarm_failure_{graph_bs}"), + ) + continue + try: + torch.cuda.synchronize() + except Exception as exc: + # An asynchronous kernel failure may poison the context; + # continuing with the tensor fallback is not safe. + raise RuntimeError( + f"DSpark fused scheduler failed its startup CUDA gate for G={graph_bs}" + ) from exc + ready.append(graph_bs) + + self._dspark_fused_schedule_ready_sizes = set(ready) + logger.info(f"DSpark fused scheduler prewarm complete for graph batch sizes {ready}") + + @torch.inference_mode() + def _warmup_dspark_ragged_compressor_metadata(self) -> None: + """Compile DSv4 ragged metadata for a non-CUDA-graph batch size. + + Whole-model CUDA-graph warmup exercises only configured graph batch + sizes. The first mixed attention-DP iteration can contain an + intermediate number of generation requests, which otherwise triggers + a long ``torch.compile`` specialization immediately before a TP + collective. Ranks without generation work reach the collective first + and can trip the executor hang detector while their peers compile. + + One non-bucket probe forces the dynamic specialization during engine + warmup, before the hang detector starts. It uses the metadata's real + persistent buffers so the tensor shape and stride guards match serving. + Those buffers are created under inference mode, so all probe and + cleanup mutations must run under the same mode. + """ + if not self._dspark_confidence_enabled or not self._dspark_trims_submitted_tokens: + return + + metadata = getattr(self, "attn_metadata", None) + required_attributes = ( + "_compute_compressed_mask", + "_compute_gen_compressed_position_ids", + "_compress_ratios_sorted", + "compressed_mask_cuda", + "compressed_position_ids_cuda", + "cu_new_comp_kv_cuda", + "max_draft_tokens", + "new_comp_kv_lens_cuda", + "past_kv_lens_cuda", + ) + missing_attributes = [ + name for name in required_attributes if metadata is None or not hasattr(metadata, name) + ] + if missing_attributes: + raise RuntimeError( + "DSpark ragged metadata warmup requires the DSv4 compressor " + "contract; missing " + ", ".join(missing_attributes) + ) + + graph_batch_sizes = {int(batch_size) for batch_size in self._config.cuda_graph_batch_sizes} + num_generations = next( + ( + candidate + for candidate in range(int(self._config.max_batch_size) - 1, 0, -1) + if candidate not in graph_batch_sizes + ), + None, + ) + if num_generations is None: + return + + compress_ratios = [int(ratio) for ratio in metadata._compress_ratios_sorted] + tokens_per_generation = 1 + int(metadata.max_draft_tokens) + total_compressed_tokens = {} + gen_output_offsets = {ratio: 0 for ratio in compress_ratios} + prepared_ratios = [] + try: + for ratio in compress_ratios: + compressed_tokens_per_generation = (tokens_per_generation + ratio - 1) // ratio + total_tokens = num_generations * compressed_tokens_per_generation + past_kv_lens = metadata.past_kv_lens_cuda[ratio] + cu_new_comp = metadata.cu_new_comp_kv_cuda[ratio] + new_comp = metadata.new_comp_kv_lens_cuda[ratio] + compressed_positions = metadata.compressed_position_ids_cuda[ratio] + compressed_mask = metadata.compressed_mask_cuda[ratio] + if ( + past_kv_lens.numel() < num_generations + or cu_new_comp.numel() < num_generations + 1 + or new_comp.numel() < num_generations + or compressed_positions.numel() < total_tokens + or compressed_mask.numel() < total_tokens + ): + raise RuntimeError( + "DSpark ragged compressor metadata buffers cannot " + f"hold the compile probe with {num_generations} " + "generation requests and compression ratio " + f"{ratio}" + ) + + total_compressed_tokens[ratio] = total_tokens + prepared_ratios.append(ratio) + past_kv_lens[:num_generations].zero_() + new_comp[:num_generations].fill_(compressed_tokens_per_generation) + cu_new_comp[: num_generations + 1].copy_( + torch.arange( + num_generations + 1, dtype=cu_new_comp.dtype, device=cu_new_comp.device + ) + ) + cu_new_comp[: num_generations + 1].mul_(compressed_tokens_per_generation) + + logger.info( + "DSpark ragged metadata warmup: compiling one non-graph " + f"shape with {num_generations} generation requests" + ) + metadata._compute_gen_compressed_position_ids( + metadata.past_kv_lens_cuda, + metadata.cu_new_comp_kv_cuda, + metadata.compressed_position_ids_cuda, + 0, + num_generations, + tokens_per_generation, + compress_ratios, + gen_output_offsets, + ) + metadata._compute_compressed_mask( + metadata.new_comp_kv_lens_cuda, + metadata.cu_new_comp_kv_cuda, + metadata.compressed_mask_cuda, + num_generations, + total_compressed_tokens, + compress_ratios, + ) + torch.cuda.synchronize() + finally: + for ratio in prepared_ratios: + total_tokens = total_compressed_tokens[ratio] + metadata.past_kv_lens_cuda[ratio][:num_generations].zero_() + metadata.new_comp_kv_lens_cuda[ratio][:num_generations].zero_() + metadata.cu_new_comp_kv_cuda[ratio][: num_generations + 1].zero_() + metadata.compressed_position_ids_cuda[ratio][:total_tokens].zero_() + metadata.compressed_mask_cuda[ratio][:total_tokens].zero_() + torch.cuda.synchronize() + logger.info("DSpark ragged metadata warmup complete") + + def _set_warmup_ragged_windows(self, batch, verify_bucket: int, draft_len: int) -> None: + """Give a capture batch per-request windows summing to ``verify_bucket``. + + The warmup batch must hit the bucket total exactly, or the graph is + keyed under a total no runtime batch produces and every ragged step + silently falls back to eager. Any split with the right sum captures + the right shape (the raggedness lives in the contents of qo_indptr). + """ + requests = batch.generation_requests + if not requests: + return + from tensorrt_llm._torch.speculative.ragged_helpers import RaggedVerifyLayout + + lens = torch.ones(len(requests), dtype=torch.int32) + layout = RaggedVerifyLayout.from_verify_lens( + lens, graph_num_tokens=int(verify_bucket), total_verify_tokens=len(requests) + ) + filled = layout.fill_bucket(max_verify_len=1 + int(draft_len)) + for request, tokens in zip(requests, filled.verify_lens.tolist()): + request.py_verify_len = int(tokens) - 1 + # Publish to the runner too: the graph key carries this value, and the + # runtime fit never runs during capture. + self.cuda_graph_runner.agreed_ragged_bucket = int(verify_bucket) + + def ragged_verify_token_buckets(self, padded_bs: int) -> List[int]: + """Measured compact ``V`` cells captured for one exact ``G``.""" + if self._config.spec_config is None: + return [] + from tensorrt_llm._torch.speculative.dspark_planner import ExactSpsCostTable + + exact_table = getattr(self, "_dspark_sps_cost_table", None) + if not isinstance(exact_table, ExactSpsCostTable): + raise RuntimeError( + "DSpark ragged verification requires an authenticated exact SPS table" + ) + return list(exact_table.production_candidate_budgets(int(padded_bs))) + + def _get_spec_worker(self): + """Return the one-model speculative worker when it is initialized.""" + return getattr(self.model, "spec_worker", None) + + def fit_ragged_verify_lens( + self, + generation_requests, + verify_lens: List[int], + exact_shape: Tuple[int, int, int], + peer_stats: Optional[List[List[int]]] = None, + exact_zero_real: bool = False, + ) -> Optional[int]: + """Validate and publish a globally selected exact ``(G,V)`` shape. + + ``peer_stats`` is every attention-DP rank's ``[num_requests, + total_tokens, all_can_graph]``, including this rank's: the bucket must + be sized from the cross-rank maximum or ranks pick different shapes, + and ``all_can_graph`` is the group's answer to whether that maximum + will be taken at all. + + ``exact_shape`` is the agreed ``(G,V,pad_tokens)`` from the exact SPS + selector. Changing V here would execute a different measured cell from + the one the policy priced, so this method only validates the split. + + ``exact_zero_real`` means this rank contributes no logical requests + but carries the one scheduler-owned attention-DP dummy. The dummy is + mapped onto the high row when ``V % G != 0``; CUDA padding supplies + the remaining low/high rows without publishing output or KV ownership + for a real request. + """ + runner = self.cuda_graph_runner + # Cleared on entry so every early return leaves no stale value: a + # bucket carried over from a previous step would key into the wrong + # graph. The padded row count follows the same rule. + runner.agreed_ragged_bucket = None + runner.ragged_zero_real_high_rows = 0 + self._dspark_last_padded_bs = None + if not runner.enabled or not verify_lens: + return None + if len(verify_lens) != len(generation_requests): + # A partially windowed batch would go ragged in the token layout + # while the spec metadata stayed uniform; refuse instead. + logger.debug( + f"DSpark ragged: got {len(verify_lens)} verify lengths for " + f"{len(generation_requests)} generation requests; falling back " + f"to uniform scheduling" + ) + return None + max_verify_len = 1 + int(self._config.spec_config.max_draft_len) + token_lens = [1 + int(v) for v in verify_lens] + + from tensorrt_llm._torch.speculative.ragged_helpers import resolve_ragged_pad_split + + # Resolve rows first, then tokens: the bucket grid depends on the + # widest rank's row count. `all_can_graph` (third peer-stat element; + # absent means single-rank) says whether _get_padded_batch will take + # the cross-rank row maximum at all -- only when EVERY rank can graph. + all_can_graph = bool(peer_stats[0][2]) if peer_stats and len(peer_stats[0]) > 2 else True + if not all_can_graph: + # No rank will replay a graph this step, and _get_padded_batch + # pads differently than this fit assumes on such steps, so the + # budget would go to rows that never appear. Decline (conservative: + # an eager ragged step would still save compute). + logger.debug( + "DSpark ragged: some rank cannot run a CUDA graph this step, " + "so no captured bucket applies; falling back to uniform" + ) + return None + widest_rows = ( + max([int(s[0]) for s in peer_stats], default=len(token_lens)) + if peer_stats + else len(token_lens) + ) + padded_bs = int(exact_shape[0]) + expected_padded_bs = runner._round_up_batch_size(widest_rows) + if padded_bs != expected_padded_bs: + logger.warning( + "DSpark exact ragged shape disagrees with graph ladder: " + f"selected G={padded_bs}, expected G={expected_padded_bs}" + ) + return None + if padded_bs == 0: + return None + # The fit assumes the batch will actually be padded to `padded_bs`; + # padding may decline, and then the fitted total lands in no captured + # bucket. Decline to go ragged rather than fit against rows that will + # not exist (`will_pad_to` mirrors only the SIZE guards; the group's + # graph answer is `all_can_graph` above). + if not runner.will_pad_to(padded_bs, len(token_lens)): + logger.debug( + f"DSpark ragged: padding to {padded_bs} rows is not available " + f"for {len(token_lens)} requests, so the bucket grid derived " + f"from it would not be realised; falling back to uniform" + ) + return None + buckets = self.ragged_verify_token_buckets(padded_bs) + if not buckets: + return None + bucket = int(exact_shape[1]) + if bucket not in buckets: + logger.warning(f"DSpark exact V={bucket} was not captured for G={padded_bs}") + return None + + full_bucket = int(padded_bs) * int(max_verify_len) + if bucket == full_bucket: + # A full verifier budget has no ragged work reduction. Preserve + # the native static-K path and its ordinary multi-G tail handling. + return None + + if exact_zero_real: + if ( + len(generation_requests) != 1 + or not generation_requests[0].is_attention_dp_dummy + or len(verify_lens) != 1 + ): + raise RuntimeError( + "DSpark zero-real exact fit lost its single scheduled " + "attention-DP dummy invariant" + ) + quotient, remainder = divmod(int(bucket), int(padded_bs)) + if quotient < 1 or quotient + int(remainder > 0) > max_verify_len: + raise RuntimeError( + "DSpark zero-real exact bucket has no bounded low/high " + f"row split: G={padded_bs}, V={bucket}, " + f"max_tokens={max_verify_len}" + ) + expected_scheduled_window = quotient - 1 + int(remainder > 0) + if int(verify_lens[0]) != expected_scheduled_window: + raise RuntimeError( + "DSpark zero-real scheduled dummy window differs from " + f"the exact quotient/remainder split: got={verify_lens[0]}, " + f"expected={expected_scheduled_window}" + ) + if remainder > 1: + draft_len = int(self._config.spec_config.max_draft_len) + if draft_len not in runner.secondary_padding_dummy_requests: + raise RuntimeError( + "DSpark secondary padding dummy disappeared after " + "the all-rank exact-cell agreement" + ) + self._dspark_last_padded_bs = int(padded_bs) + self._dspark_last_num_real = 0 + runner.agreed_ragged_bucket = int(bucket) + runner.ragged_pad_verify_len = int(quotient) - 1 + runner.ragged_zero_real_high_rows = int(remainder) + generation_requests[0].py_verify_len = expected_scheduled_window + return int(bucket) + + # Pad rows are appended later as a *single shared dummy object*, so + # they all carry the same window and their contribution must be + # decided here, not left to the fill. They take the minimum (one + # token) so real requests get the slack, growing only when the real + # rows cannot absorb the bucket. + n_real = len(token_lens) + n_pad = padded_bs - n_real + floor_tokens = sum(token_lens) + fixed_pad_len = int(exact_shape[2]) + split = resolve_ragged_pad_split( + bucket=int(bucket), + num_real_requests=n_real, + total_real_tokens=floor_tokens, + padded_bs=int(padded_bs), + max_verify_len=max_verify_len, + fixed_pad_len=fixed_pad_len, + ) + if split is None: + logger.warning( + f"DSpark ragged: bucket {bucket} admits no pad-row window " + f"for {n_real} real requests and {n_pad} pad rows; falling " + f"back to uniform scheduling" + ) + return None + pad_len = split.pad_len + real_target = split.real_target + if real_target != sum(token_lens): + logger.warning( + "DSpark exact layout changed between policy and fit: " + f"real tokens={sum(token_lens)}, target={real_target}, " + f"G={padded_bs}, V={bucket}, pad={pad_len}" + ) + return None + # The exact selector already returned a bounded split whose real rows + # sum to the measured cell; repacking it cannot change a value. + published = token_lens + + # Published only now, past the last way this fit can fail: a stale + # published bucket on a fallback step is the state behind the ragged + # IMA. The row count and pad window feed the fresh-confidence device + # prologue; `ragged_pad_verify_len` is also stamped on the shared dummy + # by _get_padded_batch. + self._dspark_last_padded_bs = int(padded_bs) + self._dspark_last_num_real = int(n_real) + runner.agreed_ragged_bucket = int(bucket) + runner.ragged_pad_verify_len = pad_len - 1 + + for request, tokens in zip(generation_requests, published): + request.py_verify_len = int(tokens) - 1 + return int(bucket) + + def _select_dspark_windows_with_fused_fallback( + self, + *, + select_fn: Callable[..., Any], + selector_kwargs: dict[str, Any], + padded_bs: int, + ) -> Any: + """Run the optional scheduler once, retiring only a failed graph G.""" + from tensorrt_llm._torch.speculative.dspark_schedule import DSparkFusedScheduleError + + ready_sizes = getattr(self, "_dspark_fused_schedule_ready_sizes", set()) + use_fused_exact = padded_bs in ready_sizes + try: + return select_fn(**selector_kwargs, use_fused_exact=use_fused_exact) + except DSparkFusedScheduleError as exc: + if not use_fused_exact: + raise + self._retire_dspark_fused_shape(padded_bs, exc, "scheduler") + return select_fn(**selector_kwargs, use_fused_exact=False) + + def _retire_dspark_fused_shape(self, padded_bs: int, exc: Exception, component: str) -> None: + """Permanently route one failed graph-row shape to the tensor oracle.""" + ready_sizes = set(getattr(self, "_dspark_fused_schedule_ready_sizes", set())) + ready_sizes.discard(padded_bs) + self._dspark_fused_schedule_ready_sizes = ready_sizes + failure_counts = getattr(self, "_dspark_fused_schedule_failure_counts", {}) + failure_counts[padded_bs] = failure_counts.get(padded_bs, 0) + 1 + self._dspark_fused_schedule_failure_counts = failure_counts + logger.warning_once( + f"DSpark fused {component} failed for G={padded_bs}; disabling " + "that shape and rerunning the exact tensor fallback. " + f"{type(exc).__name__}: {exc}", + key=f"dspark_fused_{component}_runtime_failure_{padded_bs}", + ) + + def _apply_device_window_prologue(self, inputs, new_tensors_device) -> bool: + """Re-rank this step's verify windows on device, with fresh confidence. + + Runs after ``_prepare_inputs`` and before the graph replay, all device + ops on the current stream -- stream order makes every write visible to + the replayed graph, and nothing here reads device data back to the + host. The host has already agreed ``(padded_bs, bucket)`` and staged a + SHAPE SPLIT of the (lagged) budget through the normal fit; this + prologue re-distributes the same real/pad token totals by the verified + block's OWN confidence (ranking lag zero) and overwrites the layout's + content: verify lens, qo_indptr, the per-token row maps, the packed + input/position/draft tokens, and the per-request kv-length delta from + the shape split to the true windows. + + Returns True when applied; False when a precondition fails (the step + then runs the shape split as-is, which is a valid window assignment). + """ + from tensorrt_llm._torch.speculative.dspark_device_select import ( + gather_packed_draft_tokens, + materialize_compact_layout, + select_windows_device, + ) + from tensorrt_llm._torch.speculative.dspark_schedule import DSparkFusedScheduleError + + runner = self.cuda_graph_runner + budget = self._dspark_device_budget + self._dspark_device_budget = None + if self._dspark_host_window_step: + return False + if budget is None or runner.agreed_ragged_bucket is None: + return False + if not getattr(self, "_dspark_prev_covers_batch", False): + # The input gathers below address new_tokens_device by + # previous_batch_indices in batch order; a request without a + # previous device tensor breaks that addressing. + return False + worker = self._get_spec_worker() + planner = worker.verify_planner + if planner is None or worker.staged_confidence_buffer() is None: + return False + if worker.batch_slot_view(1) is None: + return False + spec_metadata = inputs.get("spec_metadata") + attn_metadata = inputs.get("attn_metadata") + if spec_metadata is None or attn_metadata is None: + return False + if self._config.use_mrope: + # Host-selected ragged windows build the complete three-axis + # MRoPE positions correctly. The device prologue currently owns + # only the scalar position buffer, so decline fresh reranking + # rather than leave the 3-D positions on the stale shape split. + return False + apply_device_layout = getattr(attn_metadata, "apply_device_ragged_layout", None) + if not callable(apply_device_layout): + # Device-selected row ownership is an attention-backend + # capability. Decline before mutating any shared buffers when a + # backend only supports the ordinary host-selected ragged layout. + return False + if int(attn_metadata.num_contexts) != 0: + return False + + bucket = int(runner.agreed_ragged_bucket) + padded_bs = int(self._dspark_last_padded_bs) + n_real = int(self._dspark_last_num_real) + pad_len_tok = int(runner.ragged_pad_verify_len) + 1 + real_tokens = bucket - (padded_bs - n_real) * pad_len_tok + cfg = planner.cfg + # The scheduler cannot grant more than the real rows can absorb under + # the published split; the fill tops up any shortfall. + budget = max(0, min(int(budget), real_tokens - n_real * (cfg.min_verify_len + 1))) + + lens_buf = self.ragged_verify_lens_cuda + qo_buf = self.ragged_qo_indptr_cuda + + expected_stamp = worker.verified_draft_seq_cuda() + stamps = worker.confidence_stamp_buffer() if expected_stamp is not None else None + selector_kwargs = dict( + confidence_logits=worker.staged_confidence_buffer(), + slot_idx=worker.batch_slot_view(padded_bs), + num_real=n_real, + budget=budget, + graph_num_tokens=bucket, + cfg=cfg, + apply_calibration=planner.apply_calibration, + stamp=stamps, + expected_stamp=expected_stamp, + pad_len=pad_len_tok, + workspace=getattr(self, "_dspark_device_window_workspace", None), + ) + result = self._select_dspark_windows_with_fused_fallback( + select_fn=select_windows_device, + selector_kwargs=selector_kwargs, + padded_bs=padded_bs, + ) + + fused_snapshot_valid = False + if result.workspace is not None: + try: + materialize_compact_layout( + result=result, + old_verify_lens=lens_buf[:padded_bs], + old_qo_indptr=qo_buf[: padded_bs + 1], + batch_slots=self.previous_batch_indices_cuda[:n_real], + new_tokens=new_tensors_device.new_tokens, + new_tokens_lens=new_tensors_device.new_tokens_lens, + next_draft_tokens=new_tensors_device.next_draft_tokens, + input_ids=self.input_ids_cuda[:bucket], + position_ids=self.position_ids_cuda[:bucket], + previous_pos_indices=self.previous_pos_indices_cuda[:real_tokens], + previous_pos_id_offsets=self.previous_pos_id_offsets_cuda[:real_tokens], + draft_tokens=self.draft_tokens_cuda[: real_tokens - n_real], + kv_lens=attn_metadata.kv_lens_cuda[:padded_bs], + previous_kv_lens_offsets=self.previous_kv_lens_offsets_cuda[:padded_bs], + num_real=n_real, + real_tokens=real_tokens, + ) + except DSparkFusedScheduleError as exc: + # Row state commits in the final kernel, so a synchronous + # launch failure before it is safe to rerun from the original + # shape split. Every token output is overwritten below. + self._retire_dspark_fused_shape(padded_bs, exc, "layout materializer") + fused_snapshot_valid = bool(getattr(exc, "past_seen_valid", False)) + result = select_windows_device(**selector_kwargs, use_fused_exact=False) + else: + req_idx = result.req_idx + spec_metadata.remap_expanded_sampling_params(req_idx, bucket) + apply_device_layout(result.verify_lens, req_idx, result.kv_correction) + return True + + # Tensor oracle/fallback: keep the previous independently implemented + # materialization unchanged. It also repairs every token output if an + # optional fused launch failed synchronously above. + split_lens = lens_buf[:padded_bs].clone() + split_qo = qo_buf[: padded_bs + 1].to(torch.long) + if fused_snapshot_valid: + # The first fused launch saved these values before any token write. + past_seen = self._dspark_device_window_workspace.past_seen[:padded_bs] + else: + past_seen = self.position_ids_cuda[split_qo[:-1]].clone() + + lens_buf[:padded_bs].copy_(result.verify_lens) + qo_buf[: padded_bs + 1].copy_(result.qo_indptr) + # Host staging includes the window in both num_cached and seq_lens_kv: + # kv_lens = past + 2*shape_lens. Fresh windows therefore require twice + # the row delta here and its negative in the replay overlap offset. + window_delta = result.verify_lens - split_lens + attn_metadata.kv_lens_cuda[:padded_bs] += 2 * window_delta + # After replay adds its offset, the sum is past + fresh_window + new_len. + self.previous_kv_lens_offsets_cuda[:padded_bs] -= window_delta.to( + self.previous_kv_lens_offsets_cuda.dtype + ) + + req_idx = result.req_idx + spec_metadata.remap_expanded_sampling_params(req_idx, bucket) + device = req_idx.device + flat = torch.arange(bucket, device=device) + offset = flat - result.qo_indptr.to(torch.long)[req_idx] + prev_slots = self.previous_batch_indices_cuda[:n_real].to(torch.long) + slots_tok = prev_slots[req_idx.clamp(max=n_real - 1)] + + new_tokens_device = new_tensors_device.new_tokens + new_tokens_lens_device = new_tensors_device.new_tokens_lens + next_draft_tokens_device = new_tensors_device.next_draft_tokens + + self.input_ids_cuda[:bucket] = ( + new_tokens_device.transpose(0, 1)[slots_tok, offset] + .flatten() + .to(self.input_ids_cuda.dtype) + ) + self.position_ids_cuda[:bucket] = past_seen[req_idx] + offset.to(past_seen.dtype) + # Overlap corrections gather by each token's OWNER; rebuild both the + # index and the per-token offset it feeds (mirrors the host staging at + # the previous_pos_indices block). + self.previous_pos_indices_cuda[:real_tokens] = slots_tok[:real_tokens].to( + self.previous_pos_indices_cuda.dtype + ) + self.previous_pos_id_offsets_cuda[:real_tokens].copy_( + new_tokens_lens_device[slots_tok[:real_tokens]] + ) + # Draft-token storage excludes every bonus/anchor position, including + # the last row's anchor when the allocation is full. + real_draft = real_tokens - n_real + if real_draft > 0: + self.draft_tokens_cuda[:real_draft].copy_( + gather_packed_draft_tokens( + next_draft_tokens=next_draft_tokens_device, + batch_slots=prev_slots, + verify_lens=result.verify_lens, + qo_indptr=result.qo_indptr, + num_real=n_real, + total_draft_tokens=real_draft, + ).to(self.draft_tokens_cuda.dtype) + ) + + apply_device_layout(result.verify_lens, req_idx, result.kv_correction) + return True + + @staticmethod + def _ragged_token_lens(generation_requests) -> Optional[List[int]]: + """Each generation request's token window, or None if the batch is uniform. + + ``py_verify_len`` counts drafted positions; the token window adds the + bonus position. None unless *every* request carries a window -- + partially windowed batches are silent token misattribution. + """ + verify_lens = [getattr(request, "py_verify_len", None) for request in generation_requests] + if not verify_lens or any(v is None for v in verify_lens): + return None + return [1 + int(v) for v in verify_lens] + + def _publish_gen_token_layout( + self, + attn_metadata, + generation_requests, + *, + runtime_draft_len: int, + host_authoritative: bool = False, + ) -> None: + """Hand the attention metadata this step's gen-token layout, before it + prepares. + + Split out of :meth:`_attach_ragged_verify_layout` for ordering: + ``attn_metadata.prepare()`` is the only consumer of + ``ragged_verify_lens`` and runs well before the spec metadata is + assembled. The uniform stride is published here too, because + ``max_draft_tokens`` is a static buffer-sizing ceiling that does not + move when a shorter tier is chosen. It cannot live in + ``update_spec_dec_param`` either: that runs against the base metadata, + and ``prepare()`` runs against the per-key CUDA-graph copy. + """ + if attn_metadata is None: + return + if hasattr(attn_metadata, "runtime_tokens_per_gen_step"): + attn_metadata.runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step( + runtime_draft_len + ) + if hasattr(attn_metadata, "ragged_verify_lens"): + attn_metadata.ragged_verify_lens = ( + self._ragged_token_lens(generation_requests) + if self._dspark_trims_submitted_tokens + else None + ) + if hasattr(attn_metadata, "device_windows_mode"): + # Tells the attention prepare that the host window VALUES are a + # shape split (bounds only); the true windows land on device + # through apply_device_ragged_layout after prepare. + attn_metadata.device_windows_mode = ( + self._dspark_trims_submitted_tokens + and self._dspark_device_windows + and not host_authoritative + ) + + def _pinned_host(self, key: str, values, dtype) -> torch.Tensor: + """A persistent pinned staging buffer holding ``values``. + + Async H2D sources must outlive the queued copy -- PyTorch does not + extend the lifetime of an async source, so a temporary tensor can be + reclaimed while the DMA still reads it. Buffers are keyed by name and + grown monotonically. Two slots support the token and draft gathers + issued under the same key in one preparation pass. The executor calls + :meth:`wait_for_input_copy` before the next pass mutates host input, so + no additional per-slot event or synchronization is needed. + """ + values = list(values) + slot = 1 - self._pinned_host_active.get(key, 1) + bufs = self._pinned_host_cache.setdefault(key, [None, None]) + buf = bufs[slot] + if buf is None or buf.numel() < len(values) or buf.dtype != dtype: + buf = torch.empty(max(len(values), 1), dtype=dtype, pin_memory=prefer_pinned()) + bufs[slot] = buf + self._pinned_host_active[key] = slot + view = buf[: len(values)] + if values: + view.copy_(torch.tensor(values, dtype=dtype)) + return view + + def _attach_ragged_verify_layout( + self, spec_metadata, attn_metadata, generation_requests + ) -> None: + """Publish this step's per-request verify windows to both metadatas. + + Only the DSpark confidence scheduler sets ``py_verify_len``; every + other path leaves the fields None. Acceptance needs the windows to + slice correctly; the DSA indexer needs them to expand kv_lens/block + tables per query token. + """ + if not self._dspark_trims_submitted_tokens: + spec_metadata.verify_lens = None + spec_metadata.qo_indptr = None + spec_metadata.total_verify_tokens = None + if attn_metadata is not None and hasattr(attn_metadata, "ragged_verify_lens"): + attn_metadata.ragged_verify_lens = None + return + + token_lens = self._ragged_token_lens(generation_requests) + if token_lens is None: + spec_metadata.verify_lens = None + spec_metadata.qo_indptr = None + spec_metadata.total_verify_tokens = None + if attn_metadata is not None and hasattr(attn_metadata, "ragged_verify_lens"): + attn_metadata.ragged_verify_lens = None + return + + from tensorrt_llm._torch.speculative.ragged_helpers import build_qo_indptr + + n = len(token_lens) + # Persistent buffers: a captured graph baked in the address it saw at + # capture time, so a fresh tensor would be invisible to every replay. + lens_view = self.ragged_verify_lens_cuda[:n] + lens_view.copy_( + self._pinned_host("ragged_verify_lens", token_lens, torch.int32), non_blocking=True + ) + indptr_view = self.ragged_qo_indptr_cuda[: n + 1] + indptr_view.copy_(build_qo_indptr(lens_view), non_blocking=True) + spec_metadata.verify_lens = lens_view + spec_metadata.qo_indptr = indptr_view + spec_metadata.total_verify_tokens = sum(token_lens) + if attn_metadata is not None and hasattr(attn_metadata, "ragged_verify_lens"): + attn_metadata.ragged_verify_lens = token_lens + + def _ragged_gather_indices( + self, slots: List[int], counts: List[int] + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Row/column index pair for gathering a ragged block out of a + [num_slots, max_width] device tensor. + + ``slots[i]`` contributes ``counts[i]`` entries taken from columns + ``0..counts[i]-1``, concatenated in batch order. This replaces the + ``tensor[slots, :width]`` strided gather used when every request + verifies the same number of positions. + """ + rows, cols = ragged_gather_index_lists(slots, counts) + rows_dev = self._pinned_host("gather_rows", rows, torch.long).to("cuda", non_blocking=True) + cols_dev = self._pinned_host("gather_cols", cols, torch.long).to("cuda", non_blocking=True) + return (rows_dev, cols_dev) + def _can_use_steady_gen_fast_prepare( self, scheduled_requests: ScheduledRequests, @@ -3545,6 +4513,9 @@ def _helix_pack_extend(request, group: int) -> int: # will contain previous batch indices of generation requests previous_batch_indices = [] previous_pos_indices = [] + previous_batch_token_start = None + previous_batch_tokens_per_request = [] if self._dspark_trims_submitted_tokens else None + extend_tokens_emitted = 0 runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step(runtime_draft_len) runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 for request in extend_requests: @@ -3564,6 +4535,11 @@ def _helix_pack_extend(request, group: int) -> int: or request.is_dummy or request.py_batch_idx is None ): + num_draft_tokens = get_draft_token_length(request) + if self._dspark_trims_submitted_tokens: + num_draft_tokens = ( + get_request_tokens_per_gen_step(request, 1 + num_draft_tokens) - 1 + ) # get token ids, including input token ids and draft token ids. For these dummy requests, # no need to copy the token ids. if not (request.is_attention_dp_dummy or request.is_cuda_graph_dummy): @@ -3571,10 +4547,12 @@ def _helix_pack_extend(request, group: int) -> int: input_ids.append(request.get_tokens(0)[request.context_current_position]) else: input_ids.append(request.get_last_tokens(0)) - input_ids.extend(request.py_draft_tokens) - draft_tokens.extend(request.py_draft_tokens) + submitted_draft_tokens = request.py_draft_tokens + if self._dspark_trims_submitted_tokens: + submitted_draft_tokens = submitted_draft_tokens[:num_draft_tokens] + input_ids.extend(submitted_draft_tokens) + draft_tokens.extend(submitted_draft_tokens) # get other ids and lengths - num_draft_tokens = get_draft_token_length(request) past_seen_token_num = ( request.context_current_position if is_promoted_context @@ -3603,6 +4581,9 @@ def _helix_pack_extend(request, group: int) -> int: num_cached_tokens_per_seq.append( past_seen_token_num - request.py_num_compressed_tokens ) + if self._dspark_trims_submitted_tokens: + request.cached_tokens = past_seen_token_num + extend_tokens_emitted += 1 + num_draft_tokens if _has_cp_helix: # Verify group [base, base+group) in GLOBAL positions. # On a helix gen worker the request's token list is the @@ -3626,30 +4607,36 @@ def _helix_pack_extend(request, group: int) -> int: previous_batch_idx = request.py_batch_idx request.py_batch_idx = request.py_seq_slot - sequence_lengths.append(runtime_tokens_per_gen_step) + request_tokens_per_step = runtime_tokens_per_gen_step + if self._dspark_trims_submitted_tokens: + request_tokens_per_step = get_request_tokens_per_gen_step( + request, runtime_tokens_per_gen_step + ) + sequence_lengths.append(request_tokens_per_step) num_accepted_draft_tokens.append(request.py_num_accepted_draft_tokens) past_seen_token_num = request.max_beam_num_tokens - 1 - draft_lens.append(runtime_draft_token_buffer_width) + draft_lens.append(request_tokens_per_step - 1) gather_ids.extend( - list(range(len(position_ids), len(position_ids) + runtime_tokens_per_gen_step)) + list(range(len(position_ids), len(position_ids) + request_tokens_per_step)) ) position_ids.extend( - list( - range( - past_seen_token_num, past_seen_token_num + runtime_tokens_per_gen_step - ) - ) + list(range(past_seen_token_num, past_seen_token_num + request_tokens_per_step)) ) # previous tensor previous_batch_indices.append(previous_batch_idx) - previous_pos_indices.extend([previous_batch_idx] * runtime_tokens_per_gen_step) + previous_pos_indices.extend([previous_batch_idx] * request_tokens_per_step) + if previous_batch_token_start is None: + previous_batch_token_start = extend_tokens_emitted + if previous_batch_tokens_per_request is not None: + previous_batch_tokens_per_request.append(request_tokens_per_step) + extend_tokens_emitted += request_tokens_per_step num_cached_tokens_per_seq.append( - past_seen_token_num - + runtime_tokens_per_gen_step - - request.py_num_compressed_tokens + past_seen_token_num + request_tokens_per_step - request.py_num_compressed_tokens ) + if self._dspark_trims_submitted_tokens: + request.cached_tokens = past_seen_token_num + request_tokens_per_step if _has_cp_helix: # In-flight predecessor: mirror the non-helix convention # above -- positions are packed from the stale base (the @@ -3658,7 +4645,7 @@ def _helix_pack_extend(request, group: int) -> int: # in recompute_helix_spec_buffers overrides them). The # base is reconstructed GLOBALLY (see the no-previous # branch: the token list is rank-local under helix). - group = runtime_tokens_per_gen_step + group = request_tokens_per_step base = _helix_pack_extend(request, group) local_full = _helix_local_len_host(base + group) helix_owned_new_tokens.append(0) @@ -3668,7 +4655,7 @@ def _helix_pack_extend(request, group: int) -> int: and spec_config.spec_dec_mode.extend_ctx(self._config.attention_backend) and spec_config.is_linear_tree ): - prompt_lengths.append(runtime_tokens_per_gen_step) + prompt_lengths.append(request_tokens_per_step) else: prompt_lengths.append(request.py_prompt_len) @@ -3917,6 +4904,9 @@ def _helix_pack_extend(request, group: int) -> int: mrope_delta_read_seq_slots.clear() previous_batch_len = len(previous_batch_indices) + self._dspark_prev_covers_batch = previous_batch_len > 0 and previous_batch_len == len( + extend_requests + ) - len(extend_dummy_requests) def previous_seq_slots_device(): previous_batch_indices_host = torch.tensor( @@ -3960,18 +4950,43 @@ def previous_seq_slots_device(): if previous_batch_len > 0: previous_slots = previous_seq_slots_device() + is_ragged_gen = previous_batch_tokens_per_request is not None and any( + tokens != runtime_tokens_per_gen_step + for tokens in previous_batch_tokens_per_request + ) # previous input ids - previous_batch_tokens = previous_batch_len * runtime_tokens_per_gen_step - new_tokens = new_tokens_device.transpose(0, 1)[ - previous_slots, :runtime_tokens_per_gen_step - ].flatten() + previous_batch_tokens = len(previous_pos_indices) + if is_ragged_gen: + token_rows, token_cols = self._ragged_gather_indices( + previous_batch_indices, previous_batch_tokens_per_request + ) + new_tokens = new_tokens_device.transpose(0, 1)[token_rows, token_cols].flatten() + else: + new_tokens = new_tokens_device.transpose(0, 1)[ + previous_slots, :runtime_tokens_per_gen_step + ].flatten() self.input_ids_cuda[num_tokens : num_tokens + previous_batch_tokens].copy_( new_tokens, non_blocking=True ) # previous draft tokens - previous_batch_draft_tokens = previous_batch_len * runtime_draft_token_buffer_width - if runtime_draft_token_buffer_width > 0: + if is_ragged_gen: + previous_batch_draft_tokens = previous_batch_tokens - previous_batch_len + if previous_batch_draft_tokens > 0: + draft_rows, draft_cols = self._ragged_gather_indices( + previous_batch_indices, + [tokens - 1 for tokens in previous_batch_tokens_per_request], + ) + self.draft_tokens_cuda[ + num_draft_tokens : num_draft_tokens + previous_batch_draft_tokens + ].copy_( + next_draft_tokens_device[draft_rows, draft_cols].flatten(), + non_blocking=True, + ) + elif runtime_draft_token_buffer_width > 0: + previous_batch_draft_tokens = ( + previous_batch_len * runtime_draft_token_buffer_width + ) self.draft_tokens_cuda[ num_draft_tokens : num_draft_tokens + previous_batch_draft_tokens ].copy_( @@ -3981,7 +4996,18 @@ def previous_seq_slots_device(): non_blocking=True, ) # prepare data for the preprocess inputs - kv_len_offsets_device = new_tokens_lens_device - runtime_tokens_per_gen_step + if is_ragged_gen: + tokens_per_previous_request = torch.tensor( + previous_batch_tokens_per_request, + dtype=torch.long, + pin_memory=prefer_pinned(), + ).to(new_tokens_lens_device.device, non_blocking=True) + previous_kv_len_offsets = ( + new_tokens_lens_device[previous_slots] - tokens_per_previous_request + ).clamp_(min=-(runtime_draft_len + 1), max=runtime_draft_len + 1) + else: + kv_len_offsets_device = new_tokens_lens_device - runtime_tokens_per_gen_step + previous_kv_len_offsets = kv_len_offsets_device[previous_slots] previous_pos_indices_host = torch.tensor( previous_pos_indices, dtype=torch.int, pin_memory=prefer_pinned() ) @@ -4010,10 +5036,14 @@ def previous_seq_slots_device(): # Already set to '0' during initialization. num_extend_reqeust_wo_dummy = len(extend_requests) - len(extend_dummy_requests) - self.previous_pos_id_offsets_cuda[ - (num_extend_reqeust_wo_dummy - previous_batch_len) - * runtime_tokens_per_gen_step : num_extend_reqeust_wo_dummy + pos_offsets_start = ( + previous_batch_token_start + if is_ragged_gen + else (num_extend_reqeust_wo_dummy - previous_batch_len) * runtime_tokens_per_gen_step + ) + self.previous_pos_id_offsets_cuda[ + pos_offsets_start : pos_offsets_start + previous_batch_tokens ].copy_( new_tokens_lens_device[self.previous_pos_indices_cuda[0:previous_batch_tokens]], non_blocking=True, @@ -4021,7 +5051,7 @@ def previous_seq_slots_device(): self.previous_kv_lens_offsets_cuda[ num_extend_reqeust_wo_dummy - previous_batch_len : num_extend_reqeust_wo_dummy - ].copy_(kv_len_offsets_device[previous_slots], non_blocking=True) + ].copy_(previous_kv_len_offsets, non_blocking=True) elif new_tokens_device is not None: seq_slots_device = previous_seq_slots_device() @@ -4109,7 +5139,22 @@ def previous_seq_slots_device(): helix_owned_new_tokens=(helix_owned_new_tokens if helix_spec_active else None), ) - if not attn_metadata.is_cuda_graph: + if self._dspark_trims_submitted_tokens: + self._publish_gen_token_layout( + attn_metadata, + scheduled_requests.generation_requests, + runtime_draft_len=runtime_draft_len, + host_authoritative=bool(getattr(self, "_dspark_host_window_step", False)), + ) + refresh_seq_lens = not attn_metadata.is_cuda_graph + if ( + not refresh_seq_lens + and getattr(attn_metadata, "is_ragged_verify", False) + and attn_metadata.seq_lens is not None + and len(sequence_lengths) == attn_metadata.seq_lens.shape[0] + ): + refresh_seq_lens = True + if refresh_seq_lens: # Assumes seq lens do not change between CUDA graph invocations. This applies # to draft sequences too. This means that all draft sequences must be padded. attn_metadata.seq_lens = torch.tensor( @@ -4339,6 +5384,10 @@ def previous_seq_slots_device(): ] if context_prompt_lookahead is not None: spec_metadata.populate_context_prompt_lookahead(context_prompt_lookahead) + if self._dspark_trims_submitted_tokens: + self._attach_ragged_verify_layout( + spec_metadata, attn_metadata, scheduled_requests.generation_requests + ) # No-op for non 1-model spec_metadata.populate_sampling_params_for_one_model(scheduled_requests.all_requests()) spec_metadata.prepare() @@ -4383,6 +5432,13 @@ def previous_seq_slots_device(): ) if not is_dummy: + self.previous_request_ids = all_gen_request_ids + if self._dspark_trims_submitted_tokens: + self.previous_verify_lens = [ + getattr(request, "py_verify_len", None) + for request in scheduled_requests.generation_requests + ] + self.has_previous_device_draft = next_draft_tokens_device is not None # Record the steady-state generation cache when this pass handled # purely non-dummy generation requests that all carried a previous # overlap-scheduler tensor (previous_batch_len == _n_gen implies @@ -4642,6 +5698,7 @@ def _forward_decoder( if spec_metadata is not None: spec_metadata.update_is_all_greedy_sample(padded_graph_requests.all_requests()) self._sync_group_all_greedy_sample(spec_metadata) + self._validate_dspark_draft_sampling(spec_metadata) peft_cache_data_type = None if self._lora.cuda_graph_manager is not None: @@ -4745,6 +5802,13 @@ def _forward_decoder( self._prepare_inputs_event = torch.cuda.Event() self._prepare_inputs_event.record() + if ( + can_run_graph + and not is_dummy + and getattr(self, "_dspark_device_budget", None) is not None + ): + self._apply_device_window_prologue(inputs, new_tensors_device) + breakable_runner = self.breakable_cuda_graph_runner with with_shared_pool(self.cuda_graph_runner.get_graph_pool()): @@ -4988,6 +6052,16 @@ def _execute_logit_post_processors(self, scheduled_requests: ScheduledRequests, logits_processors = getattr(request, "py_logits_post_processors", None) if logits_processors: + if ( + not is_context_request + and getattr(request, "py_verify_len", None) is not None + ): + raise RuntimeError( + "per-request logits post-processors are not supported with DSpark " + "ragged verification: the processor interface assumes one uniform " + "row stride and would modify another request's packed logits after " + "a short verify window" + ) token_ids = ( [request.get_tokens(0)] if is_context_request diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b9e44d5ad892..4b2cb1df265d 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -830,6 +830,68 @@ def cuda_graph_runner(self) -> Optional[CUDAGraphRunner]: return None return self._runner.cuda_graph_runner + @property + def _dspark_confidence_enabled(self) -> bool: + return isinstance( + self._runner, + DecoderRunner) and self._runner._dspark_confidence_enabled + + @property + def _dspark_trims_submitted_tokens(self) -> bool: + return isinstance( + self._runner, + DecoderRunner) and self._runner._dspark_trims_submitted_tokens + + @property + def _dspark_sps_cost_table(self): + if not isinstance(self._runner, DecoderRunner): + return None + return self._runner._dspark_sps_cost_table + + @property + def _dspark_exact_candidate_cells(self): + if not isinstance(self._runner, DecoderRunner): + return () + return self._runner._dspark_exact_candidate_cells + + @property + def _dspark_exact_identity_words(self): + if not isinstance(self._runner, DecoderRunner): + return (0, ) * 8 + return self._runner._dspark_exact_identity_words + + @property + def _dspark_device_budget(self): + if not isinstance(self._runner, DecoderRunner): + return None + return getattr(self._runner, "_dspark_device_budget", None) + + @property + def _dspark_mixed_runtime(self): + """Optional mixed controller owned by the actual decoder runner.""" + if not isinstance(self._runner, DecoderRunner): + return None + return self._runner._dspark_mixed_runtime + + @_dspark_device_budget.setter + def _dspark_device_budget(self, value) -> None: + if isinstance(self._runner, DecoderRunner): + self._runner._dspark_device_budget = value + + def ragged_verify_token_buckets(self, batch_size: int): + if not isinstance(self._runner, DecoderRunner): + return () + return self._runner.ragged_verify_token_buckets(batch_size) + + def fit_ragged_verify_lens(self, *args, **kwargs): + assert isinstance(self._runner, DecoderRunner) + return self._runner.fit_ragged_verify_lens(*args, **kwargs) + + def _get_spec_worker(self): + if isinstance(self._runner, DecoderRunner): + return self._runner._get_spec_worker() + return None + @property def metrics(self) -> dict[str, float]: """Return model-engine warmup time metrics.""" diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 2a7cdd83d5e3..e277088802f7 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -14,8 +14,8 @@ from contextlib import contextmanager, nullcontext from enum import IntEnum from queue import Queue -from typing import (TYPE_CHECKING, Dict, Iterable, Iterator, List, Literal, - Optional, Tuple, Union) +from typing import (TYPE_CHECKING, Callable, Dict, Iterable, Iterator, List, + Literal, Optional, Tuple, Union) import torch from strenum import StrEnum @@ -79,6 +79,7 @@ from .adp_iter_stats import ADPIterStatsBuffer from .connectors.kv_cache_connector import KvCacheConnectorManager from .connectors.kv_cache_layout import build_kv_cache_layout_v2 +from .cuda_graph_runner import ADPShapeAgreement from .disagg_adapter import PyExecutorEffects, PyExecutorRequestRegistry from .dwdp import DwdpManager from .error_classification import ErrorBudget @@ -129,6 +130,150 @@ _UNBOUNDED_STATS_MAX_LEN = -1 _V2_RETENTION_IGNORED_LOG_KEY = "kv_cache_manager_v2_retention_config_ignored" +_DSPARK_EXACT_WIRE_PREFIX_LEN = 16 +_DSPARK_ADP_WIRE_TRAILER_LEN = 3 +_DSPARK_EXACT_NATIVE_YIELD_INDEX = _DSPARK_EXACT_WIRE_PREFIX_LEN - 1 + + +def _encode_dspark_exact_expected_yields( + *, exact_local, num_cells: int, + yield_scale: int) -> Tuple[int, Tuple[int, ...]]: + """Encode one rank's exact yields without splitting the collective.""" + if yield_scale <= 0: + raise ValueError("DSpark exact-policy yield scale must be positive") + num_cells = int(num_cells) + if num_cells < 0: + raise ValueError("DSpark exact-policy cell count must be nonnegative") + if exact_local is None: + return 0, tuple(-1 for _ in range(num_cells)) + + native = float(exact_local.native_expected_yield) + native_valid = math.isfinite(native) and native >= 0.0 + native_wire = int(math.ceil(native * yield_scale)) if native_valid else 0 + local_values = tuple(exact_local.compact_expected_yields) + compact_wires = [] + for cell_index in range(num_cells): + value = (float(local_values[cell_index]) + if cell_index < len(local_values) else math.nan) + if (not native_valid or not math.isfinite(value) or value < 0.0 + or value > native): + compact_wires.append(-1) + else: + compact_wires.append(int(math.floor(value * yield_scale))) + return native_wire, tuple(compact_wires) + + +def _decode_dspark_exact_expected_yields( + *, + payloads: List[List[int]], + exact_cells: Iterable[Tuple[int, int]], + graph_batch_size: int, + yield_scale: int, +) -> Tuple[float, Dict[int, float], Dict[int, float]]: + """Decode aggregate yields and worst-rank loss per real request.""" + if yield_scale <= 0: + raise ValueError("DSpark exact-policy yield scale must be positive") + native_expected_yield = sum( + float(payload[_DSPARK_EXACT_NATIVE_YIELD_INDEX]) + for payload in payloads) / yield_scale + compact_expected_yields = {} + compact_max_yield_losses_per_request = {} + for cell_index, (cell_graph_batch_size, + verifier_budget) in enumerate(exact_cells): + if cell_graph_batch_size != graph_batch_size: + continue + rank_yields = [ + float(payload[_DSPARK_EXACT_WIRE_PREFIX_LEN + cell_index]) + for payload in payloads + ] + infeasible = any(value < 0.0 for value in rank_yields) + compact_expected_yields[verifier_budget] = (0.0 if infeasible else + sum(rank_yields) / + yield_scale) + active_rank_losses = [] + if not infeasible: + for payload, compact_yield_wire in zip(payloads, rank_yields): + num_real_requests = int(payload[1]) + if num_real_requests <= 0: + continue + native_yield_wire = float( + payload[_DSPARK_EXACT_NATIVE_YIELD_INDEX]) + if compact_yield_wire > native_yield_wire: + raise RuntimeError( + "DSpark exact-policy compact yield exceeds native yield" + ) + active_rank_losses.append( + (native_yield_wire - compact_yield_wire) / yield_scale / + num_real_requests) + compact_max_yield_losses_per_request[verifier_budget] = max( + active_rank_losses, default=0.0) + return ( + native_expected_yield, + compact_expected_yields, + compact_max_yield_losses_per_request, + ) + + +def _validate_dspark_exact_bucket(*, exact_shape, bucket) -> None: + """Require the graph key to match the globally selected exact cell.""" + if exact_shape is None: + return + expected_bucket = int(exact_shape[1]) + if bucket is None or int(bucket) != expected_bucket: + raise RuntimeError( + "DSpark exact policy did not publish its selected captured " + f"verifier bucket V={expected_bucket}; got {bucket}") + + +def _validate_dspark_adp_acceptance_gate( + *, confidence_enabled: bool, attention_dp_enabled: bool, + speculation_gate_enabled: bool) -> None: + """Reject a rank-local gate that can split the ADP collective protocol.""" + if confidence_enabled and attention_dp_enabled and speculation_gate_enabled: + raise ValueError( + "DSpark confidence scheduling with attention DP does not support " + "acceptance_rate_window_size/acceptance_rate_threshold because " + "the per-rank gate can break collective ordering") + + +def _classify_dspark_exact_generation_rows( + generation_requests: List[LlmRequest], +) -> Tuple[Optional[List[LlmRequest]], bool]: + """Return logical exact rows and whether one idle ADP dummy is scheduled.""" + if not generation_requests: + return None, False + real_requests = [ + request for request in generation_requests if not request.is_dummy + ] + if len(real_requests) == len(generation_requests): + return real_requests, False + one_adp_dummy = bool(not real_requests and len(generation_requests) == 1 + and generation_requests[0].is_attention_dp_dummy) + if one_adp_dummy: + return real_requests, True + return None, False + + +def _dspark_exact_common_graph_batch_size( + payloads: List[List[int]], round_up: Callable[[int], int]) -> int: + """Resolve a common exact G, preserving zero as all-rank termination.""" + widest_rows = max((int(payload[1]) for payload in payloads), default=0) + return 0 if widest_rows <= 0 else int(round_up(widest_rows)) + + +def _dspark_exact_secondary_padding_ready( + payloads: List[List[int]], + *, + graph_batch_size: int, + verifier_budget: int, + secondary_ready_index: int, +) -> bool: + """Whether every zero-real peer can represent the selected remainder.""" + if verifier_budget % graph_batch_size <= 1: + return True + return all( + int(payload[1]) != 0 or bool(int(payload[secondary_ready_index])) + for payload in payloads) class _ADPForwardIntent(IntEnum): @@ -611,6 +756,12 @@ def __init__( threshold = getattr(spec_config, 'acceptance_rate_threshold', None) if window and threshold is not None: self.speculation_gate = SpeculationGate(window, threshold) + _validate_dspark_adp_acceptance_gate( + confidence_enabled=getattr(self.model_engine, + "_dspark_confidence_enabled", False), + attention_dp_enabled=self.enable_attention_dp, + speculation_gate_enabled=self.speculation_gate is not None, + ) # response used data self.response_lock = threading.Lock() @@ -1961,7 +2112,7 @@ def get_latest_iteration_stats(self): Returns the per-iterations statistics computed since last call to this method. Contains at most iter_stats_max_iterations iterations. """ - if self.enable_iter_perf_stats == False: + if not self.enable_iter_perf_stats: return [] latest_stats = (IterationStats(), None) @@ -3774,13 +3925,414 @@ def _drain_relay_sends_before_forward(self, microbatch_id: int) -> None: self.dist.pp_size) % self.num_micro_batches self.wait_on_pp_send_handles(self.send_handles, due_microbatch_id) + def _stage_dspark_confidence(self, worker, planner, gen_requests, + mixed) -> None: + buffer = worker.staged_confidence_buffer() + if buffer is None: + return + worker.bump_draft_seq() + if mixed is not None and mixed.costs is not None: + from .dspark_mixed_runtime import stage_mixed_ownership + + planner.stage_confidence(buffer, worker.confidence_stamp_buffer()) + stage_mixed_ownership(planner, worker, int(self.iter_counter), + gen_requests) + else: + planner.stage_confidence(buffer) + + def _dspark_confidence_draft_len(self, + scheduled_batch: ScheduledRequests) -> int: + """Agree one DSpark verify shape and stage confidence asynchronously.""" + max_draft_len = self.model_engine.max_draft_len + worker = self.model_engine._get_spec_worker() + planner = getattr(worker, "verify_planner", None) + if planner is None: + return max_draft_len + mixed = self.model_engine._dspark_mixed_runtime + if mixed is not None: + mixed.reset_step() + + runtime_cost_table = getattr(self.model_engine, + "_dspark_sps_cost_table", None) + if runtime_cost_table is not None and planner.cost_table is not runtime_cost_table: + planner.install_runtime_cost_table(runtime_cost_table) + + gen_requests = scheduled_batch.generation_requests + is_distributed = self.dist is not None and self.dist.tp_size > 1 + compute_windows = bool( + getattr(self.model_engine, "_dspark_trims_submitted_tokens", False)) + + ragged_lens = None + device_budget = None + exact_local = None + exact_table = planner.exact_cost_table + if exact_table is None: + raise RuntimeError( + "DSpark confidence scheduling has no authenticated exact SPS table" + ) + exact_cells = tuple( + getattr(self.model_engine, "_dspark_exact_candidate_cells", ())) + if not exact_cells: + exact_cells = exact_table.candidate_cells() + exact_requests: Optional[List[LlmRequest]] = gen_requests + exact_zero_real_adp_dummy = False + exact_requests, exact_zero_real_adp_dummy = ( + _classify_dspark_exact_generation_rows(gen_requests)) + + # A stale capacity must never re-rank a later batch. + self.model_engine._dspark_device_budget = None + if compute_windows: + if scheduled_batch.context_requests: + pass + else: + if exact_requests is not None: + exact_rows = [ + worker.confidence_row_for(request.py_request_id) + for request in exact_requests + ] + exact_local = planner.prepare_exact_sps_decision( + num_gen_requests=len(exact_requests), rows=exact_rows) + + local_rows = (exact_local.num_requests if exact_local is not None else + len(ragged_lens) if ragged_lens is not None else 0) + local_tokens = (sum( + 1 + int(value) + for value in ragged_lens) if ragged_lens is not None else 0) + local_payload = [ + 1 if exact_local is not None or ragged_lens is not None else 0, + int(local_rows), + int(local_tokens), + 1 if scheduled_batch.can_run_cuda_graph else 0, + ] + + exact_yield_scale = 1_000_000 + exact_identity_words = (0, ) * 8 + exact_identity_words = tuple( + getattr( + self.model_engine, + "_dspark_exact_identity_words", + exact_table.collective_identity_words, + )) + native_yield_wire, compact_yields_wire = _encode_dspark_exact_expected_yields( + exact_local=exact_local, + num_cells=len(exact_cells), + yield_scale=exact_yield_scale, + ) + local_payload.extend([ + 1, + 1 if exact_local is not None else 0, + len(exact_cells), + *exact_identity_words, + native_yield_wire, + *compact_yields_wire, + ]) + + runner = self.model_engine.cuda_graph_runner + local_payload.extend([ + int(scheduled_batch.batch_size), + 1 if planner.max_verify_len in getattr( + runner, "padding_dummy_requests", {}) else 0, + 1 if planner.max_verify_len in getattr( + runner, "secondary_padding_dummy_requests", {}) else 0, + ]) + if mixed is not None: + local_payload.extend( + mixed.local_vote( + scheduled_batch, + worker, + planner, + int(self.iter_counter), + int(self.dist.tp_rank) if is_distributed else 0, + executor_eligible=bool( + compute_windows + and self._dspark_adp_shape_cache_enabled()), + )) + payloads = (self.dist.tp_allgather_int64(local_payload).tolist() + if is_distributed else [local_payload]) + + exact_payload_end = _DSPARK_EXACT_WIRE_PREFIX_LEN + len(exact_cells) + expected_payload_len = exact_payload_end + _DSPARK_ADP_WIRE_TRAILER_LEN + mixed_peers = None + if mixed is not None: + payloads, mixed_peers = mixed.split_votes(payloads, + expected_payload_len) + if any(len(payload) != expected_payload_len for payload in payloads): + raise RuntimeError( + "DSpark ADP agreement payload has an inconsistent shape") + + exact_table_flags = [int(payload[4]) for payload in payloads] + if len(set(exact_table_flags)) != 1: + raise RuntimeError( + "DSpark ranks disagree on whether an exact SPS table is installed" + ) + if exact_table_flags[0]: + exact_cell_counts = [int(payload[6]) for payload in payloads] + exact_identities = [ + tuple(int(value) for value in payload[7:15]) + for payload in payloads + ] + if (len(set(exact_cell_counts)) != 1 + or exact_cell_counts[0] != len(exact_cells) + or len(set(exact_identities)) != 1): + raise RuntimeError( + "DSpark exact SPS ranks disagree on the authenticated " + "(G,V), T(G,V), or minimum-gain table identity") + + all_can_graph = all(int(payload[3]) for payload in payloads) + peer_stats = None + if not all(int(payload[0]) for payload in payloads): + ragged_lens = None + exact_local = None + else: + peer_stats = [[ + int(payload[1]), + int(payload[2]), 1 if all_can_graph else 0 + ] for payload in payloads] + + batch_size_index = exact_payload_end + padding_dummy_index = exact_payload_end + 1 + secondary_padding_dummy_index = exact_payload_end + 2 + peer_batch_sizes = tuple( + int(payload[batch_size_index]) for payload in payloads) + widest_batch_size = max(peer_batch_sizes, default=0) + common_padded_batch_size = runner._round_up_batch_size_with_draft_len( + widest_batch_size, planner.max_verify_len) + padding_needed = any( + int(payload[batch_size_index]) != common_padded_batch_size + for payload in payloads) + padding_supported = bool( + not padding_needed + or (runner.padding_enabled + and common_padded_batch_size <= runner.config.batch_size)) + padding_ready = bool( + common_padded_batch_size > 0 and padding_supported and all( + int(payload[batch_size_index]) == common_padded_batch_size + or int(payload[padding_dummy_index]) for payload in payloads)) + exact_policy_agreed = bool(all(int(payload[5]) for payload in payloads)) + agreement = ADPShapeAgreement( + iteration=int(self.iter_counter), + batch_identity=id(scheduled_batch), + local_batch_size=int(scheduled_batch.batch_size), + peer_batch_sizes=peer_batch_sizes, + all_can_graph=all_can_graph, + widest_batch_size=widest_batch_size, + graph_batch_size=common_padded_batch_size, + draft_len=int(planner.max_verify_len), + padding_ready=padding_ready, + ) + runner.adp_shape_agreement = agreement + if (mixed_peers is not None and agreement.can_queue + and mixed.apply(scheduled_batch, worker, planner, mixed_peers, + int(self.iter_counter))): + # Common context-aware get_padding_params supplies the BREAKABLE + # body. A decode-only peer must not replay its local full graph. + runner.adp_shape_agreement = dataclasses.replace( + agreement, + all_can_graph=False, + ragged_bucket=None, + reuse_graph_shape=False, + ) + self._stage_dspark_confidence(worker, planner, gen_requests, mixed) + return int(planner.max_verify_len) + if not agreement.can_queue: + runner.agreed_ragged_bucket = None + runner.ragged_pad_verify_len = 0 + runner.ragged_zero_real_high_rows = 0 + for request in gen_requests: + request.py_verify_len = None + return int(planner.max_verify_len) + + exact_shape = None + if exact_local is not None: + if not all_can_graph: + exact_local = None + else: + if any(not int(payload[5]) for payload in payloads): + raise RuntimeError( + "DSpark exact SPS collective payloads disagree on the " + "measured (G,V) grid") + common_graph_batch_size = _dspark_exact_common_graph_batch_size( + payloads, runner._round_up_batch_size) + if common_graph_batch_size == 0: + exact_local = None + else: + from ..speculative.dspark_planner import \ + select_exact_sps_candidate + + ( + native_expected_yield, + compact_expected_yields, + compact_max_yield_losses_per_request, + ) = _decode_dspark_exact_expected_yields( + payloads=payloads, + exact_cells=exact_cells, + graph_batch_size=common_graph_batch_size, + yield_scale=exact_yield_scale, + ) + selected_verifier_budget = select_exact_sps_candidate( + graph_batch_size=common_graph_batch_size, + native_expected_yield=native_expected_yield, + compact_expected_yields=compact_expected_yields, + compact_max_yield_losses_per_request=( + compact_max_yield_losses_per_request), + cost_table=exact_table, + ) + if selected_verifier_budget == 0 or not ( + _dspark_exact_secondary_padding_ready( + payloads, + graph_batch_size=common_graph_batch_size, + verifier_budget=selected_verifier_budget, + secondary_ready_index= + secondary_padding_dummy_index, + )): + exact_local = None + else: + ragged_lens, local_exact_budget, exact_pad_tokens = ( + planner.allocate_exact_sps_candidate( + exact_local, + graph_batch_size=common_graph_batch_size, + verifier_budget=selected_verifier_budget, + )) + if exact_local.num_requests == 0: + if not exact_zero_real_adp_dummy: + raise RuntimeError( + "zero-real exact SPS allocation has no " + "authenticated attention-DP dummy row") + _, remainder = divmod(selected_verifier_budget, + common_graph_batch_size) + scheduled_tokens = int(exact_pad_tokens) + int( + remainder > 0) + ragged_lens = [scheduled_tokens - 1] + exact_shape = ( + common_graph_batch_size, + selected_verifier_budget, + exact_pad_tokens, + ) + if planner.device_windows and local_exact_budget > 0: + device_budget = int(local_exact_budget) + + from ..speculative.dspark_verify import \ + _exact_cell_geometry + + peer_stats = [] + for payload in payloads: + peer_rows = int(payload[1]) + peer_geometry = _exact_cell_geometry( + num_real=peer_rows, + graph_batch_size=common_graph_batch_size, + verifier_budget=selected_verifier_budget, + min_verify_len=planner.cfg.min_verify_len, + max_verify_len=planner.max_verify_len, + ) + if peer_geometry is None: + raise RuntimeError( + "Selected exact SPS cell is infeasible for " + f"peer rows={peer_rows}, " + f"G={common_graph_batch_size}, " + f"V={selected_verifier_budget}") + _, peer_real_tokens, _ = peer_geometry + peer_stats.append([ + peer_rows, peer_real_tokens, + 1 if all_can_graph else 0 + ]) + + bucket = None + ragged_active = False + if ragged_lens is not None: + if exact_shape is None: + raise RuntimeError( + "DSpark exact allocation produced no measured graph shape") + if all_can_graph: + bucket = self.model_engine.fit_ragged_verify_lens( + gen_requests, + ragged_lens, + peer_stats=peer_stats, + exact_shape=exact_shape, + exact_zero_real=exact_zero_real_adp_dummy, + ) + ragged_active = bucket is not None + if ragged_active and device_budget is not None: + self.model_engine._dspark_device_budget = int(device_budget) + else: + runner.agreed_ragged_bucket = None + runner.ragged_pad_verify_len = 0 + runner.ragged_zero_real_high_rows = 0 + for request, window in zip(gen_requests, ragged_lens): + request.py_verify_len = int(window) + ragged_active = True + + reuse_graph_shape = bool(exact_policy_agreed and runner.enabled + and all_can_graph and padding_ready + and not runner.config.use_mrope) + _validate_dspark_exact_bucket(exact_shape=exact_shape, bucket=bucket) + if (reuse_graph_shape and exact_shape is not None + and exact_shape[0] != common_padded_batch_size): + raise RuntimeError( + "DSpark exact policy selected a graph batch size that " + "disagrees with the cached ADP padding shape") + runner.adp_shape_agreement = dataclasses.replace( + agreement, + ragged_bucket=None if bucket is None else int(bucket), + reuse_graph_shape=reuse_graph_shape, + ) + + if not ragged_active: + for request in gen_requests: + request.py_verify_len = None + runner.agreed_ragged_bucket = None + runner.ragged_pad_verify_len = 0 + runner.ragged_zero_real_high_rows = 0 + + self._stage_dspark_confidence(worker, planner, gen_requests, mixed) + return int(planner.max_verify_len) + def _handle_dynamic_draft_len(self, scheduled_batch: ScheduledRequests) -> None: - """Synchronize draft length and buffers before preparing resources.""" - update_draft_len(self.model_engine, - scheduled_batch, - speculation_permanently_disabled=self. - speculation_permanently_disabled) + """Synchronize draft width and confidence policy before KV admission.""" + spec_config = getattr(self.model_engine, "spec_config", None) + confidence_driven = bool( + spec_config is not None + and getattr(spec_config, "enable_confidence_scheduling", False)) + if not confidence_driven or self.speculation_permanently_disabled: + # Main owns the static, scheduled and disabled draft-buffer rules. + update_draft_len(self.model_engine, + scheduled_batch, + speculation_permanently_disabled=self. + speculation_permanently_disabled) + return + + signature = ( + int(self.iter_counter), + id(scheduled_batch), + int(scheduled_batch.batch_size), + int(scheduled_batch.num_context_requests), + int(scheduled_batch.num_generation_requests), + ) + if getattr(self, "_dspark_dynamic_handled_signature", + None) == signature: + return + runtime_draft_len = self._dspark_confidence_draft_len(scheduled_batch) + agreement = self.model_engine.cuda_graph_runner.adp_shape_agreement + if (agreement is not None and agreement.matches( + scheduled_batch, self.iter_counter, padded=False) + and not agreement.can_queue): + self.model_engine.runtime_draft_len = runtime_draft_len + else: + update_draft_len(self.model_engine, + scheduled_batch, + draft_len=runtime_draft_len) + self._dspark_dynamic_handled_signature = signature + + def _dspark_adp_shape_cache_enabled(self) -> bool: + runner = getattr(self.model_engine, "cuda_graph_runner", None) + return bool( + self.enable_attention_dp and self.dist is not None + and self.dist.tp_size > 1 and runner is not None + and getattr(self.model_engine, "_dspark_confidence_enabled", False) + and getattr(self, "kv_connector_manager", None) is None + and getattr(self, "kv_cache_transceiver", None) is None + and getattr(self, "drafter", None) is None + and getattr(self, "_pp_rebalance_drain_iters", None) is None) def _initialize_remote_tail_adp(self) -> None: """Reserve the execution-only idle row after final KV/graph warmup.""" @@ -3856,6 +4408,28 @@ def _can_queue(self, scheduled_batch): # can_queue_this_rank is for case that the batch is not empty on this rank, but empty on other ranks # For bs == 1, we cannot pad dummy request to make the batch non-empty since it will cause the batch size to be 2. # 1 for dummy request, 1 for the yet-to-complete but not-yet-updated request. + runner = getattr(self.model_engine, "cuda_graph_runner", None) + if self._dspark_adp_shape_cache_enabled(): + signature = ( + int(self.iter_counter), + id(scheduled_batch), + int(scheduled_batch.batch_size), + int(scheduled_batch.num_context_requests), + int(scheduled_batch.num_generation_requests), + ) + if getattr(self, "_dspark_dynamic_handled_signature", + None) != signature: + runner.adp_shape_agreement = None + self._handle_dynamic_draft_len(scheduled_batch) + agreement = runner.adp_shape_agreement + if agreement is not None and agreement.matches( + scheduled_batch, self.iter_counter, padded=False): + return agreement.can_queue, agreement.can_queue_this_rank + elif runner is not None: + # Optional surfaces can mutate the request set after policy + # selection, so never reuse their preceding agreement. + runner.adp_shape_agreement = None + if getattr(self, "_enable_remote_tail_adp", False): # An idle execution row needs no active-request or sampler slot. # Still veto an all-empty fleet (e.g. cancellations after the vote). @@ -5437,7 +6011,7 @@ def _maybe_finish_pp_rebalance(self) -> None: self._pp_rebalance_drain_iters = None def _suspend_padding_dummies_for_rebalance( - self, mgr: KVCacheManagerV2) -> List[Tuple[int, LlmRequest]]: + self, mgr: KVCacheManagerV2) -> List[Tuple[int, int, LlmRequest]]: """Suspend the CUDA-graph padding dummies before ``adjust()``. ``CUDAGraphRunner`` retains one padding dummy per captured draft @@ -5454,21 +6028,27 @@ def _suspend_padding_dummies_for_rebalance( and silently drop padded batches to eager mode for the rest of the process lifetime. - Returns the ``(draft_len, dummy)`` pairs that were suspended. + Returns ``(draft_len, variant, dummy)`` triples. Variant zero is the + low-row dummy and variant one is the high-row dummy. """ runner = getattr(self.model_engine, "cuda_graph_runner", None) if runner is None: return [] - suspended: List[Tuple[int, LlmRequest]] = [] - for draft_len, dummy in runner.padding_dummy_requests.items(): - if mgr.is_request_active(dummy.py_request_id): - mgr.suspend_request(dummy) - suspended.append((draft_len, dummy)) + suspended: List[Tuple[int, int, LlmRequest]] = [] + caches = ( + getattr(runner, "padding_dummy_requests", {}), + getattr(runner, "secondary_padding_dummy_requests", {}), + ) + for variant, cache in enumerate(caches): + for draft_len, dummy in cache.items(): + if mgr.is_request_active(dummy.py_request_id): + mgr.suspend_request(dummy) + suspended.append((draft_len, variant, dummy)) return suspended def _resume_padding_dummies_after_rebalance( self, mgr: KVCacheManagerV2, - suspended: List[Tuple[int, LlmRequest]]) -> None: + suspended: List[Tuple[int, int, LlmRequest]]) -> None: """Resume the padding dummies suspended for ``adjust()``. A real request that fails to resume is left suspended for the @@ -5486,7 +6066,8 @@ def _resume_padding_dummies_after_rebalance( runner = getattr(self.model_engine, "cuda_graph_runner", None) if runner is None: return - for draft_len, dummy in suspended: + failed_draft_lens = set() + for draft_len, variant, dummy in suspended: if mgr.resume_request(dummy): continue if getattr(self, "_enable_remote_tail_adp", @@ -5496,8 +6077,13 @@ def _resume_padding_dummies_after_rebalance( continue logger.warning( "Could not resume the CUDA graph padding dummy request " - f"(draft_len={draft_len}) after a KV pool rebalance; " + f"(draft_len={draft_len}, variant={variant}) after a KV pool " + "rebalance; " "releasing it and falling back to lazy re-creation.") + failed_draft_lens.add(draft_len) + # Resume both siblings before releasing a low/high pair. Releasing + # inline could free a resumed sibling and later try to resume it. + for draft_len in sorted(failed_draft_lens): runner.release_padding_dummy(self.resource_manager, draft_len) @contextmanager @@ -6181,6 +6767,10 @@ def _configure_csa2_remote_tail(self, request: LlmRequest) -> None: request.py_csa2_remote_tail_split = split def _validate_request(self, request: LlmRequest): + if getattr(self.model_engine, "_dspark_confidence_enabled", False): + from ..speculative.dspark import validate_confidence_request + + validate_confidence_request(self.model_engine.spec_config, request) self._configure_csa2_remote_tail(request) if (isinstance(self.kv_cache_manager, KVCacheManagerV2) and request.kv_cache_retention_config is not None): @@ -8861,7 +9451,7 @@ def _enqueue_responses(self, responses: Iterable[Tuple[int, LlmResponse]]): # (TODO: joyang) There are other types of responses, we need to sort out. if type( resp - ) == LlmResponse and req_id in self.result_wait_queues and self.result_wait_queues[ + ) is LlmResponse and req_id in self.result_wait_queues and self.result_wait_queues[ req_id] is not None: self.result_wait_queues[req_id].put_response.remote( resp.client_id, resp) diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index f1273926d683..4c740b7b831b 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -20,6 +20,7 @@ # ``DSv4DSparkDraftModel.forward`` rather than via mask-token cross-attention. import json +import math from collections import deque from dataclasses import dataclass, field from typing import TYPE_CHECKING, List, Optional @@ -37,6 +38,7 @@ if TYPE_CHECKING: from ...llmapi.llm_args import DSparkDecodingConfig + from ..pyexecutor.llm_request import LlmRequest # Unscored-slot fill: sigmoid saturates to 1.0, so an unwritten row schedules # as "verify the whole block" (fail-safe). @@ -60,7 +62,10 @@ def load_sts_temperatures_from_path(path: str) -> list[float]: ) if not isinstance(temps, list) or not temps: raise ValueError(f"{path} must contain a non-empty temperature list") - return [float(t) for t in temps] + temperatures = [float(t) for t in temps] + if any(not math.isfinite(value) or value <= 0.0 for value in temperatures): + raise ValueError("STS temperatures must be a finite positive numeric vector") + return temperatures def resolve_confidence_head(draft_model: torch.nn.Module) -> Optional[torch.nn.Module]: @@ -72,6 +77,35 @@ def resolve_confidence_head(draft_model: torch.nn.Module) -> Optional[torch.nn.M return getattr(stages[-1], "confidence_head", None) +def validate_confidence_request(spec_config: "DSparkDecodingConfig", request: "LlmRequest") -> None: + """Reject unsupported request controls before confidence-policy admission.""" + if not spec_config.enable_confidence_scheduling: + return + if getattr(request, "py_logits_post_processors", None): + raise ValueError( + "DSpark confidence scheduling does not support per-request logits post-processors" + ) + sampling = request.sampling_config + if not spec_config.use_rejection_sampling or sampling is None: + return + from tensorrt_llm.sampling_params import SamplingParams + + # Match SpecMetadata's actual draft-sampler classification; temperature + # alone is insufficient (for example, top_k=1 remains greedy). + greedy = SamplingParams.params_imply_greedy_decoding( + temperature=sampling.temperature, + top_k=sampling.top_k, + top_p=sampling.top_p, + min_p=sampling.min_p, + use_beam_search=False, + ) + if not greedy: + raise ValueError( + "DSpark confidence scheduling does not support advanced rejection draft sampling: " + "confidence is conditioned on the greedy proposal prefix" + ) + + def _dspark_position_ceiling(max_ctx: int, block_size: int, max_draft_len: int) -> int: """Return the number of RoPE entries needed by the DSv4 block drafter. @@ -233,6 +267,7 @@ def prepare(self): for rid in list(worker._req_to_slot.keys()): if rid not in current: slot = worker._req_to_slot.pop(rid) + worker._confidence_incarnations.pop(rid, None) worker._ctx_len[slot] = 0 worker._valid_len[slot] = 0 worker._position_initialized[slot] = False @@ -246,8 +281,8 @@ def prepare(self): # without this, all concurrent gen requests fall through to the shared # scratch row below and corrupt each other's draft window at batch # size > 1 (GitHub #16767). Context-prefix entries are left to - # ``_seed_context_windows``. CUDA-graph capture dummies and real - # generation requests acquire a persistent slot only once; prepare() + # ``_seed_context_windows``. Real generation requests acquire a + # persistent slot only once; prepare() # runs before every replay, so it must never reset a live slot. # ADP-idle (id 0) and high-ID CUDA-graph padding dummies use scratch. num_contexts = max(0, len(self.request_ids) - self.num_generations) @@ -397,6 +432,13 @@ def __init__( # gen path (set on the host in prepare(), so the captured forward indexes # the rolling windows through a tensor instead of a python dict lookup). self._req_to_slot = {} # request_id -> slot index + # Enabled only for admitted host-mixed scheduling. Slot reuse cannot + # authenticate a different request incarnation with an older snapshot. + self._mixed_confidence_tracking = bool( + getattr(spec_config, "confidence_mixed_evidence", None) is not None + ) + self._confidence_incarnations: dict[int, int] = {} + self._confidence_incarnation_sequence = 0 self._free_slots = deque() # available slot indices self._batch_to_slot: Optional[torch.Tensor] = None # [max_batch] long, cuda # Index of the throwaway "scratch" window row that absorbs padded / @@ -451,7 +493,15 @@ def bump_draft_seq(self) -> int: sequence of the last completed draft. """ prev = self._draft_seq_host - self._draft_seq_host += 1 + self._draft_seq_host = (prev + 1) & ((1 << 31) - 1) + if self._draft_seq_host == 0: + # Recycled stamps must not authenticate an ancient score. The + # wrap step sees expected=-1 and uses neutral confidence; later + # steps can consume only newly produced scores or neutral rows. + if self._confidence_stamp is not None: + self._confidence_stamp.zero_() + if self._confidence_logits is not None: + self._confidence_logits.fill_(_NEUTRAL_CONFIDENCE_LOGIT) if self._draft_seq_cuda is not None: self._draft_seq_cuda.fill_(self._draft_seq_host) return prev @@ -558,8 +608,9 @@ def _lazy_init(self, draft_model, spec_metadata, attn_metadata=None) -> None: self._scratch_slot = max_batch num_rows = max_batch + 1 - # CUDA-graph padding requests carry ids in - # ``[CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len, CUDA_GRAPH_DUMMY_REQUEST_ID]``, + # Both CUDA-graph padding variants, each with K+1 runtime widths, + # use ``[CUDA_GRAPH_DUMMY_REQUEST_ID - 2*(K+1) + 1, + # CUDA_GRAPH_DUMMY_REQUEST_ID]``, where K is the physical draft cap, # while real request ids start at ``max_batch_size`` and grow, so a simple # floor cleanly separates them. Together with ``ATTENTION_DP_DUMMY_REQUEST_ID`` # (0) these dummies must route to the scratch row (see ``prepare()``) and @@ -567,7 +618,9 @@ def _lazy_init(self, draft_model, spec_metadata, attn_metadata=None) -> None: # dspark -> cuda_graph_runner -> speculative.utils -> dspark import cycle. from ..pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID - self._graph_dummy_id_floor = CUDA_GRAPH_DUMMY_REQUEST_ID - self.max_draft_len + self._graph_dummy_id_floor = ( + CUDA_GRAPH_DUMMY_REQUEST_ID - 2 * (self.max_draft_len + 1) + 1 + ) self._kv_windows = torch.zeros( (num_rows, num_stages, self._win, head_dim), @@ -696,6 +749,7 @@ def _assign_slot(self, req_id: int, reset: bool) -> int: """Get (or refresh) the slot for a request; reset clears its window.""" if reset and req_id in self._req_to_slot: old = self._req_to_slot.pop(req_id) + self._confidence_incarnations.pop(req_id, None) self._ctx_len[old] = 0 self._valid_len[old] = 0 self._position_initialized[old] = False @@ -709,6 +763,9 @@ def _assign_slot(self, req_id: int, reset: bool) -> int: ) slot = self._free_slots.popleft() self._req_to_slot[req_id] = slot + if self._mixed_confidence_tracking: + self._confidence_incarnation_sequence += 1 + self._confidence_incarnations[req_id] = self._confidence_incarnation_sequence self._ctx_len[slot] = 0 self._valid_len[slot] = 0 self._position_initialized[slot] = False @@ -720,6 +777,10 @@ def _assign_slot(self, req_id: int, reset: bool) -> int: self._confidence_logits[slot].fill_(_NEUTRAL_CONFIDENCE_LOGIT) return self._req_to_slot[req_id] + def confidence_incarnation_for(self, req_id: int) -> int | None: + """Host ownership generation, independent of a recycled device slot.""" + return self._confidence_incarnations.get(req_id) + def _seed_context_windows( self, draft_model, diff --git a/tensorrt_llm/_torch/speculative/dspark_device_select.py b/tensorrt_llm/_torch/speculative/dspark_device_select.py new file mode 100644 index 000000000000..9eb9190c83e7 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_device_select.py @@ -0,0 +1,669 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Device-side verify-window selection for the pre-replay prologue. + +The host planner decides windows from a confidence snapshot that is one +iteration old (the batch reshuffles across the lag, and the read itself must +not sync). This module ranks the block that is about to be verified by its +OWN confidence instead: everything from the slot-indexed gather to the +per-token row maps is device-resident. The production exact-tier caller runs +it immediately before replay, after the previous step's draft wrote the +confidence buffer and before this step's verify attention reads the layout. +The tensor-control fallback remains suitable for capture by a caller whose +graph key owns all of its replay-varying controls. + +The split of responsibilities follows the paper's dual-timescale design: + +* the CAPACITY -- ``(padded_bs, bucket, budget)`` -- stays host-decided from + the lagged snapshot, because it is the CUDA-graph key and the attention-DP + agreement payload, both of which must exist before launch; +* the RANKING -- which requests win the budget -- is computed here from + fresh confidence, at zero staleness. + +Per-row staleness degrades to the neutral row (verify everything for that +request), mirroring the host planner's fail-open semantics. Exact-tier +feasibility is checked from host-known scalars before the fused launch; an +unsupported fused shape falls back to the established tensor schedule/fill. +""" + +from dataclasses import dataclass +from typing import Callable, Optional, Union + +import torch +import triton +import triton.language as tl + +from .dspark_schedule import ( + NEUTRAL_CONFIDENCE_LOGIT, + DSparkFusedScheduleError, + DSparkScheduleConfig, + compute_survival, + schedule_verify_lens_topk, + schedule_verify_lens_topk_fused_fill, +) +from .ragged_helpers import build_qo_indptr, build_row_maps_device, fill_bucket_device + +__all__ = [ + "DeviceWindowResult", + "DeviceWindowWorkspace", + "DSparkCompactLayoutError", + "gather_packed_draft_tokens", + "materialize_compact_layout", + "select_windows_device", +] + + +class DSparkCompactLayoutError(DSparkFusedScheduleError): + """Synchronous optional-layout failure with snapshot recovery state.""" + + def __init__(self, message: str, *, past_seen_valid: bool) -> None: + super().__init__(message) + self.past_seen_valid = past_seen_valid + + +@dataclass +class DeviceWindowResult: + """Everything the ragged layout consumers need, all device-resident. + + Attributes: + verify_lens: ``[padded_bs]`` int32 token windows (bonus included), + summing to exactly ``graph_num_tokens``; pad rows carry their + fill. + qo_indptr: ``[padded_bs + 1]`` int32 exclusive prefix sum. + req_idx: ``[graph_num_tokens]`` int64 owning-row per packed token. + kv_correction: ``[graph_num_tokens]`` int32; composed with a + ``kv_lens`` gather it yields each token's KV extent + (``refresh_ragged_row_kv_lens``). + """ + + verify_lens: torch.Tensor + qo_indptr: torch.Tensor + req_idx: torch.Tensor + kv_correction: torch.Tensor + workspace: Optional["DeviceWindowWorkspace"] = None + + +@dataclass +class DeviceWindowWorkspace: + """Stable-address scratch for the fused compact-layout prologue. + + A workspace is owned by one model engine and sliced to the active graph + shape. Keeping these tensors alive eliminates steady-state scheduler, + prefix-sum, row-map, and position-snapshot allocations. The tensor + fallback intentionally does not use this object, preserving it as an + independent oracle. + """ + + verify_lens: torch.Tensor + qo_indptr: torch.Tensor + req_idx: torch.Tensor + kv_correction: torch.Tensor + past_seen: torch.Tensor + + @classmethod + def allocate( + cls, + *, + max_rows: int, + max_tokens: int, + device: Union[str, torch.device], + ) -> "DeviceWindowWorkspace": + if max_rows < 0 or max_tokens < 0: + raise ValueError("workspace capacities must be non-negative") + return cls( + verify_lens=torch.empty(max_rows, dtype=torch.int32, device=device), + qo_indptr=torch.empty(max_rows + 1, dtype=torch.int32, device=device), + req_idx=torch.empty(max_tokens, dtype=torch.int64, device=device), + kv_correction=torch.empty(max_tokens, dtype=torch.int32, device=device), + past_seen=torch.empty(max_rows, dtype=torch.int32, device=device), + ) + + def validate(self, *, device: torch.device, num_rows: int, num_tokens: int) -> None: + tensors = ( + ("verify_lens", self.verify_lens, torch.int32, num_rows), + ("qo_indptr", self.qo_indptr, torch.int32, num_rows + 1), + ("req_idx", self.req_idx, torch.int64, num_tokens), + ("kv_correction", self.kv_correction, torch.int32, num_tokens), + ("past_seen", self.past_seen, torch.int32, num_rows), + ) + for name, tensor, dtype, capacity in tensors: + if tensor.device != device: + raise ValueError(f"workspace.{name} must be on {device}") + if tensor.dtype != dtype: + raise TypeError(f"workspace.{name} must have dtype {dtype}, got {tensor.dtype}") + if tensor.dim() != 1 or tensor.numel() < capacity: + raise ValueError( + f"workspace.{name} needs capacity {capacity}, got shape {tuple(tensor.shape)}" + ) + if not tensor.is_contiguous(): + raise ValueError(f"workspace.{name} must be contiguous") + + +@triton.jit +def _materialize_row_maps_kernel( + verify_lens_ptr, + qo_indptr_ptr, + req_idx_ptr, + correction_ptr, + MAX_TOKEN_LEN: tl.constexpr, +): + row = tl.program_id(0) + offsets = tl.arange(0, MAX_TOKEN_LEN) + length = tl.load(verify_lens_ptr + row) + token = tl.load(qo_indptr_ptr + row) + offsets + mask = offsets < length + tl.store(req_idx_ptr + token, row, mask=mask) + tl.store(correction_ptr + token, offsets - length + 1, mask=mask) + + +def _build_row_maps_workspace( + *, + workspace: DeviceWindowWorkspace, + num_rows: int, + graph_num_tokens: int, + max_token_len: int, +) -> None: + verify_lens = workspace.verify_lens[:num_rows] + qo_indptr = workspace.qo_indptr[: num_rows + 1] + qo_indptr[0].zero_() + if num_rows == 0: + return + torch.cumsum(verify_lens, dim=0, out=qo_indptr[1:]) + if verify_lens.is_cuda: + try: + _materialize_row_maps_kernel[(num_rows,)]( + verify_lens, + qo_indptr, + workspace.req_idx, + workspace.kv_correction, + MAX_TOKEN_LEN=triton.next_power_of_2(max_token_len), + ) + except Exception as exc: + raise DSparkFusedScheduleError("DSpark fused row-map materialization failed") from exc + return + + # CPU is a hardware-independent differential-test path. Production CUDA + # never enters the allocation-heavy tensor oracle below. + req_idx, correction = build_row_maps_device(verify_lens, graph_num_tokens=graph_num_tokens) + workspace.req_idx[:graph_num_tokens].copy_(req_idx) + workspace.kv_correction[:graph_num_tokens].copy_(correction) + + +@triton.jit(do_not_specialize=["num_rows"]) +def _snapshot_past_seen_kernel( + old_qo_ptr, + position_ids_ptr, + past_seen_ptr, + num_rows, + BLOCK: tl.constexpr, +): + rows = tl.arange(0, BLOCK) + mask = rows < num_rows + starts = tl.load(old_qo_ptr + rows, mask=mask, other=0) + past_seen = tl.load(position_ids_ptr + starts, mask=mask, other=0) + tl.store(past_seen_ptr + rows, past_seen, mask=mask) + + +@triton.jit( + do_not_specialize=[ + "new_tokens_stride_token", + "new_tokens_stride_slot", + "next_draft_stride_slot", + "next_draft_stride_token", + "num_tokens", + "num_real", + "real_tokens", + ] +) +def _materialize_compact_tokens_kernel( + req_idx_ptr, + qo_indptr_ptr, + past_seen_ptr, + batch_slots_ptr, + new_tokens_ptr, + new_tokens_lens_ptr, + next_draft_tokens_ptr, + input_ids_ptr, + position_ids_ptr, + previous_pos_indices_ptr, + previous_pos_id_offsets_ptr, + draft_tokens_ptr, + new_tokens_stride_token, + new_tokens_stride_slot, + next_draft_stride_slot, + next_draft_stride_token, + num_tokens, + num_real, + real_tokens, + BLOCK: tl.constexpr, +): + token = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + token_mask = token < num_tokens + row = tl.load(req_idx_ptr + token, mask=token_mask, other=0).to(tl.int32) + row_start = tl.load(qo_indptr_ptr + row, mask=token_mask, other=0) + offset = token - row_start + owner = tl.minimum(row, num_real - 1) + slot = tl.load(batch_slots_ptr + owner, mask=token_mask, other=0) + + new_token = tl.load( + new_tokens_ptr + offset * new_tokens_stride_token + slot * new_tokens_stride_slot, + mask=token_mask, + other=0, + ) + tl.store(input_ids_ptr + token, new_token, mask=token_mask) + past_seen = tl.load(past_seen_ptr + row, mask=token_mask, other=0) + tl.store(position_ids_ptr + token, past_seen + offset, mask=token_mask) + + real_mask = token < real_tokens + tl.store(previous_pos_indices_ptr + token, slot, mask=real_mask) + new_len = tl.load(new_tokens_lens_ptr + slot, mask=real_mask, other=0) + tl.store(previous_pos_id_offsets_ptr + token, new_len, mask=real_mask) + + draft_mask = token_mask & (row < num_real) & (offset > 0) + draft_token = tl.load( + next_draft_tokens_ptr + + slot * next_draft_stride_slot + + (offset - 1) * next_draft_stride_token, + mask=draft_mask, + other=0, + ) + draft_index = row_start - row + offset - 1 + tl.store(draft_tokens_ptr + draft_index, draft_token, mask=draft_mask) + + +@triton.jit(do_not_specialize=["num_rows"]) +def _commit_compact_rows_kernel( + new_lens_ptr, + new_qo_ptr, + old_lens_ptr, + old_qo_ptr, + kv_lens_ptr, + previous_kv_offsets_ptr, + num_rows, + BLOCK: tl.constexpr, +): + rows = tl.arange(0, BLOCK) + row_mask = rows < num_rows + new_lens = tl.load(new_lens_ptr + rows, mask=row_mask, other=0) + old_lens = tl.load(old_lens_ptr + rows, mask=row_mask, other=0) + delta = new_lens - old_lens + kv_lens = tl.load(kv_lens_ptr + rows, mask=row_mask, other=0) + previous_offsets = tl.load(previous_kv_offsets_ptr + rows, mask=row_mask, other=0) + tl.store(kv_lens_ptr + rows, kv_lens + 2 * delta, mask=row_mask) + tl.store(previous_kv_offsets_ptr + rows, previous_offsets - delta, mask=row_mask) + tl.store(old_lens_ptr + rows, new_lens, mask=row_mask) + qo_mask = rows <= num_rows + new_qo = tl.load(new_qo_ptr + rows, mask=qo_mask, other=0) + tl.store(old_qo_ptr + rows, new_qo, mask=qo_mask) + + +def gather_packed_draft_tokens( + *, + next_draft_tokens: torch.Tensor, + batch_slots: torch.Tensor, + verify_lens: torch.Tensor, + qo_indptr: torch.Tensor, + num_real: int, + total_draft_tokens: int, +) -> torch.Tensor: + """Gather the device-selected real draft rows, excluding bonus tokens. + + ``verify_lens`` and ``qo_indptr`` describe token windows that include one + bonus/anchor per request. The persistent draft buffer contains drafts + only, packed request-major. Construct exactly ``total_draft_tokens`` + owners so full-batch/full-K layouts never need a one-past-the-end discard + slot for anchors. + """ + if total_draft_tokens < 0: + raise ValueError("total_draft_tokens must be non-negative") + if total_draft_tokens == 0: + return next_draft_tokens.new_empty((0,)) + device = verify_lens.device + rows = torch.arange(num_real, device=device, dtype=torch.long) + draft_counts = verify_lens[:num_real].to(torch.long) - 1 + owners = torch.repeat_interleave(rows, draft_counts, output_size=total_draft_tokens) + # Removing one bonus from every preceding request converts the token + # prefix into a draft-only prefix. + draft_qo = qo_indptr[: num_real + 1].to(torch.long) - torch.arange( + num_real + 1, device=device, dtype=torch.long + ) + flat = torch.arange(total_draft_tokens, device=device, dtype=torch.long) + offsets = flat - draft_qo[owners] + slots = batch_slots[:num_real].to(torch.long)[owners] + return next_draft_tokens[slots, offsets] + + +def materialize_compact_layout( + *, + result: DeviceWindowResult, + old_verify_lens: torch.Tensor, + old_qo_indptr: torch.Tensor, + batch_slots: torch.Tensor, + new_tokens: torch.Tensor, + new_tokens_lens: torch.Tensor, + next_draft_tokens: torch.Tensor, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + previous_pos_indices: torch.Tensor, + previous_pos_id_offsets: torch.Tensor, + draft_tokens: torch.Tensor, + kv_lens: torch.Tensor, + previous_kv_lens_offsets: torch.Tensor, + num_real: int, + real_tokens: int, +) -> None: + """Materialize a selected compact layout without per-step temporaries. + + CUDA snapshots row positions, writes token/position/overlap/draft buffers, + and commits the row delta with three allocation-free Triton launches. A + CPU implementation mirrors the former tensor code for differential tests; + it is deliberately not a second optimized production path. + + The row commit is last. Therefore a synchronous failure in either earlier + launch leaves the staged lens/qo/KV state intact and lets the caller rerun + the established tensor fallback safely. + """ + workspace = result.workspace + if workspace is None: + raise ValueError("compact materialization requires a workspace-backed result") + num_rows = result.verify_lens.numel() + num_tokens = result.req_idx.numel() + num_real = int(num_real) + real_tokens = int(real_tokens) + if not 0 < num_real <= num_rows: + raise ValueError(f"num_real must be in [1, {num_rows}], got {num_real}") + if not 0 <= real_tokens <= num_tokens: + raise ValueError(f"real_tokens must be in [0, {num_tokens}], got {real_tokens}") + device = result.verify_lens.device + workspace.validate(device=device, num_rows=num_rows, num_tokens=num_tokens) + + one_dimensional = ( + ("old_verify_lens", old_verify_lens, num_rows), + ("old_qo_indptr", old_qo_indptr, num_rows + 1), + ("batch_slots", batch_slots, num_real), + ("new_tokens_lens", new_tokens_lens, 1), + ("input_ids", input_ids, num_tokens), + ("position_ids", position_ids, num_tokens), + ("previous_pos_indices", previous_pos_indices, real_tokens), + ("previous_pos_id_offsets", previous_pos_id_offsets, real_tokens), + ("draft_tokens", draft_tokens, real_tokens - num_real), + ("kv_lens", kv_lens, num_rows), + ("previous_kv_lens_offsets", previous_kv_lens_offsets, num_rows), + ) + for name, tensor, capacity in one_dimensional: + if tensor.device != device: + raise ValueError(f"{name} must be on {device}") + if tensor.dim() != 1 or tensor.numel() < capacity: + raise ValueError(f"{name} needs capacity {capacity}, got shape {tuple(tensor.shape)}") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + if new_tokens.device != device or new_tokens.dim() not in (2, 3): + raise ValueError("new_tokens must be a 2-D or 3-D tensor on the layout device") + if new_tokens.dim() == 3 and new_tokens.shape[2] != 1: + raise ValueError("new_tokens trailing beam dimension must be one") + if next_draft_tokens.device != device or next_draft_tokens.dim() != 2: + raise ValueError("next_draft_tokens must be a 2-D tensor on the layout device") + if old_verify_lens.dtype != torch.int32 or old_qo_indptr.dtype != torch.int32: + raise TypeError("old verify_lens and qo_indptr must have dtype torch.int32") + if position_ids.dtype != torch.int32 or workspace.past_seen.dtype != position_ids.dtype: + raise TypeError("position_ids and workspace.past_seen must have dtype torch.int32") + real_draft = real_tokens - num_real + if real_draft < 0: + raise ValueError("real_tokens must include one bonus token per real row") + + if result.verify_lens.is_cuda: + row_block = triton.next_power_of_2(num_rows + 1) + token_block = 256 + try: + _snapshot_past_seen_kernel[(1,)]( + old_qo_indptr, + position_ids, + workspace.past_seen, + num_rows, + BLOCK=row_block, + ) + except Exception as exc: + raise DSparkCompactLayoutError( + "DSpark compact-layout position snapshot failed", past_seen_valid=False + ) from exc + try: + _materialize_compact_tokens_kernel[(triton.cdiv(num_tokens, token_block),)]( + result.req_idx, + result.qo_indptr, + workspace.past_seen, + batch_slots, + new_tokens, + new_tokens_lens, + next_draft_tokens, + input_ids, + position_ids, + previous_pos_indices, + previous_pos_id_offsets, + draft_tokens, + new_tokens.stride(0), + new_tokens.stride(1), + next_draft_tokens.stride(0), + next_draft_tokens.stride(1), + num_tokens, + num_real, + real_tokens, + BLOCK=token_block, + ) + _commit_compact_rows_kernel[(1,)]( + result.verify_lens, + result.qo_indptr, + old_verify_lens, + old_qo_indptr, + kv_lens, + previous_kv_lens_offsets, + num_rows, + BLOCK=row_block, + ) + except Exception as exc: + raise DSparkCompactLayoutError( + "DSpark fused compact-layout materialization failed", past_seen_valid=True + ) from exc + return + + # Hardware-independent oracle for exact differential coverage. + split_lens = old_verify_lens[:num_rows].clone() + split_qo = old_qo_indptr[: num_rows + 1].to(torch.long) + past_seen = position_ids[split_qo[:-1]].clone() + workspace.past_seen[:num_rows].copy_(past_seen) + past_seen = workspace.past_seen[:num_rows] + req_idx = result.req_idx + flat = torch.arange(num_tokens, device=device) + offset = flat - result.qo_indptr.to(torch.long)[req_idx] + prev_slots = batch_slots[:num_real].to(torch.long) + slots_tok = prev_slots[req_idx.clamp(max=num_real - 1)] + input_ids[:num_tokens].copy_( + new_tokens.transpose(0, 1)[slots_tok, offset].flatten().to(input_ids.dtype) + ) + position_ids[:num_tokens].copy_(past_seen[req_idx] + offset.to(past_seen.dtype)) + previous_pos_indices[:real_tokens].copy_(slots_tok[:real_tokens].to(previous_pos_indices.dtype)) + previous_pos_id_offsets[:real_tokens].copy_( + new_tokens_lens[slots_tok[:real_tokens]].to(previous_pos_id_offsets.dtype) + ) + if real_draft > 0: + draft_tokens[:real_draft].copy_( + gather_packed_draft_tokens( + next_draft_tokens=next_draft_tokens, + batch_slots=prev_slots, + verify_lens=result.verify_lens, + qo_indptr=result.qo_indptr, + num_real=num_real, + total_draft_tokens=real_draft, + ).to(draft_tokens.dtype) + ) + window_delta = result.verify_lens - split_lens + kv_lens[:num_rows] += 2 * window_delta + previous_kv_lens_offsets[:num_rows] -= window_delta.to(previous_kv_lens_offsets.dtype) + old_verify_lens[:num_rows].copy_(result.verify_lens) + old_qo_indptr[: num_rows + 1].copy_(result.qo_indptr) + + +def select_windows_device( + *, + confidence_logits: torch.Tensor, + slot_idx: torch.Tensor, + num_real: Union[int, torch.Tensor], + budget: Union[int, torch.Tensor], + graph_num_tokens: int, + cfg: DSparkScheduleConfig, + apply_calibration: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, + stamp: Optional[torch.Tensor] = None, + expected_stamp: Optional[torch.Tensor] = None, + pad_len: Optional[int] = None, + use_fused_exact: bool = False, + workspace: Optional[DeviceWindowWorkspace] = None, +) -> DeviceWindowResult: + """Rank the batch by fresh confidence and pack it into the agreed bucket. + + Mirrors the host chain ``_gather_rows -> apply_calibration -> + compute_survival -> schedule_verify_lens_topk -> fill_bucket`` with the + lag removed: ``confidence_logits`` is the LIVE slot-indexed buffer the + previous draft pass scattered into, not a staged snapshot. + + Args: + confidence_logits: ``[num_slots, K]`` raw logits, slot-indexed. + slot_idx: ``[padded_bs]`` int64 buffer row per batch position; + entries at or beyond ``num_real`` are never read meaningfully + (point them at any valid row, e.g. 0). + num_real: Python integer in the production pre-replay prologue, or a + 0-d integer tensor for a device-owned caller; rows past it are + padding. + budget: Python integer in the production pre-replay prologue, or a + 0-d integer tensor paired with tensor ``num_real``; verify tokens + above the floor -- the host-decided (lagged) capacity knob. + graph_num_tokens: the captured token bucket (capture constant). + cfg: scheduling bounds; ``resolved_max_verify_len + 1`` is the + per-row token ceiling. + apply_calibration: logits -> probabilities; sigmoid when None. + stamp: optional ``[num_slots]`` last-writer stamps; with + ``expected_stamp`` (0-d), rows whose stamp mismatches fall back + to the neutral row (full survival: verify everything), the same + fail-open the host planner applies to unknown slots. + pad_len: when given, pad rows carry EXACTLY this many tokens and all + fill slack goes to real rows -- the host fit's published split, + which its pre-launch copy widths depend on (see + :func:`fill_bucket_device`). The caller must clamp ``budget`` so + the real rows can absorb ``graph_num_tokens - n_pad * pad_len``. + use_fused_exact: enable the policy-neutral Triton fast path. + Production sets this only for shapes compiled successfully during + warmup and when its independent default-off switch is enabled; + false uses the established tensor schedule/fill implementation. + workspace: reusable output storage for the fused exact path. It is + ignored by the tensor fallback so that path remains an independent + allocation-owning oracle. + + Returns: + :class:`DeviceWindowResult`; every output tensor stays device-resident. + The Python-control exact path is intended for the pre-replay prologue, + while the tensor-control fallback may be captured by a compatible + caller. + """ + if confidence_logits.dim() != 2: + raise ValueError( + f"confidence_logits must be [num_slots, K], got {tuple(confidence_logits.shape)}" + ) + device = confidence_logits.device + padded_bs = slot_idx.numel() + controls_are_tensors = isinstance(num_real, torch.Tensor) + if controls_are_tensors != isinstance(budget, torch.Tensor): + raise TypeError("num_real and budget must both be Python ints or both be tensors") + if controls_are_tensors: + if num_real.dim() != 0 or budget.dim() != 0: + raise ValueError("tensor num_real and budget must both be 0-d") + # Narrow integer controls can overflow while multiplying by K or while + # casting the verification budget; production scalar controls are ordinary + # Python ints, and a device-owned caller must use a safe width. + integer_dtypes = {torch.int32, torch.int64} + if num_real.dtype not in integer_dtypes or budget.dtype not in integer_dtypes: + raise TypeError("tensor num_real and budget must both have integer dtype") + if num_real.device != device or budget.device != device: + raise ValueError("tensor num_real and budget must share confidence_logits.device") + + selected = confidence_logits.index_select(0, slot_idx.to(torch.long)) + if stamp is not None: + if expected_stamp is None: + raise ValueError("stamp requires expected_stamp") + stale = stamp.index_select(0, slot_idx.to(torch.long)) != expected_stamp + selected = torch.where( + stale.unsqueeze(1), + torch.full_like(selected, NEUTRAL_CONFIDENCE_LOGIT), + selected, + ) + + calibrate = apply_calibration or torch.sigmoid + survival = compute_survival(calibrate(selected)) + max_token_len = int(cfg.resolved_max_verify_len) + 1 + fused_exact = ( + use_fused_exact + and pad_len is not None + and not controls_are_tensors + and padded_bs <= 256 + and cfg.survival_eps > 0.0 + ) + if fused_exact: + if workspace is not None: + workspace.validate(device=device, num_rows=padded_bs, num_tokens=graph_num_tokens) + filled = schedule_verify_lens_topk_fused_fill( + survival=survival, + budget=int(budget), + num_real=int(num_real), + pad_len=int(pad_len), + cfg=cfg, + graph_num_tokens=graph_num_tokens, + out=None if workspace is None else workspace.verify_lens, + ) + else: + # Pad rows must win no policy budget. Zero is below the configured + # numerical epsilon, preserving the established tensor scheduler + # exactly. The fused scheduler masks them internally and avoids this + # arange/where temporary in production. + is_real = torch.arange(padded_bs, device=device) < num_real + survival = torch.where(is_real.unsqueeze(1), survival, torch.zeros_like(survival)) + scheduled = schedule_verify_lens_topk(survival=survival, budget=budget, cfg=cfg) + # The scheduler counts drafted positions; the token window adds the + # bonus/anchor, matching every host fill_bucket callsite. + token_lens = scheduled + 1 + filled = fill_bucket_device( + token_lens, + num_real=num_real, + graph_num_tokens=graph_num_tokens, + max_verify_len=max_token_len, + pad_fill=pad_len, + ) + if fused_exact and workspace is not None: + _build_row_maps_workspace( + workspace=workspace, + num_rows=padded_bs, + graph_num_tokens=graph_num_tokens, + max_token_len=max_token_len, + ) + return DeviceWindowResult( + verify_lens=workspace.verify_lens[:padded_bs], + qo_indptr=workspace.qo_indptr[: padded_bs + 1], + req_idx=workspace.req_idx[:graph_num_tokens], + kv_correction=workspace.kv_correction[:graph_num_tokens], + workspace=workspace, + ) + req_idx, correction = build_row_maps_device(filled, graph_num_tokens=graph_num_tokens) + return DeviceWindowResult( + verify_lens=filled, + qo_indptr=build_qo_indptr(filled), + req_idx=req_idx, + kv_correction=correction, + ) diff --git a/tensorrt_llm/_torch/speculative/dspark_mixed_evidence.py b/tensorrt_llm/_torch/speculative/dspark_mixed_evidence.py new file mode 100644 index 000000000000..bd6ba15dc7c4 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_mixed_evidence.py @@ -0,0 +1,514 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pre-construction evidence binding for measured BREAKABLE mixed costs. + +The optional mixed configuration binds measured costs and an admission receipt +to an independently observed live fingerprint. The runtime separately checks +the active route, request ownership and geometry before publishing row lengths. +Parsing authenticates evidence bytes and checks their claims; it does not run +the qualification performed by the evidence producer. +""" + +import hashlib +import json +from dataclasses import dataclass +from pathlib import Path + +from . import dspark_planner + +_BODY_BUCKETS = (96, 128, 192, 256, 384, 512) +_ACTION = "current-host-confidence-prefix-allocation-under-shared-body-capacity-v1" +_PRICE_SEMANTICS = "rank0_forward_pricing_proxy" +_GEOMETRY_FIELDS = ( + "context_rows", + "context_chunk_squared", + "context_position_product", + "context_compressed_product", + "generation_history_sum", + "generation_history_max", + "bootstrap_generations", + "real_context_requests", + "eligible_generations", + "context_requests", + "generation_requests", + "fixed_generation_rows", +) +_MAX_BYTES = 16 * 1024 * 1024 +_MAX_CELLS = 256 + + +@dataclass(frozen=True) +class MixedRuntimeIdentity: + """Independently observed live identity, not copied from the cost artifact. + + Support is limited to K3 and the eight-rank BREAKABLE body grid. The runner + checks embedded-model, PP/CP/attention-DP and capture readiness separately; + these properties are not inferred from a cost file. + """ + + source_tree_sha256: str + runtime_snapshot_sha256: str + model_config_sha256: str + gpu_class_sha256: str + topology_sha256: str + selector_identity_sha256: str + physical_k: int = 3 + world_size: int = 8 + body_buckets: tuple[int, ...] = _BODY_BUCKETS + execution_mode: str = "BREAKABLE" + + def __post_init__(self) -> None: + for field in ( + "source_tree_sha256", + "runtime_snapshot_sha256", + "model_config_sha256", + "gpu_class_sha256", + "topology_sha256", + "selector_identity_sha256", + ): + dspark_planner._require_sha256(getattr(self, field), field=field) + if ( + type(self.physical_k) is not int + or self.physical_k != 3 + or type(self.world_size) is not int + or self.world_size != 8 + or type(self.body_buckets) is not tuple + or any(type(bucket) is not int for bucket in self.body_buckets) + or self.body_buckets != _BODY_BUCKETS + or self.execution_mode != "BREAKABLE" + ): + raise ValueError("unsupported mixed physical K, topology or BREAKABLE grid") + + def payload(self) -> dict[str, object]: + return { + **{ + field: getattr(self, field) + for field in ( + "source_tree_sha256", + "runtime_snapshot_sha256", + "model_config_sha256", + "gpu_class_sha256", + "topology_sha256", + "selector_identity_sha256", + "physical_k", + "world_size", + "execution_mode", + ) + }, + "body_buckets": list(self.body_buckets), + } + + +@dataclass(frozen=True) +class MixedGeometry: + """One exact ordered-rank host-geometry/action key; never a decode ``(G,V)``. + + The twelve rank features match the accepted host price support dimensions. + They summarize host geometry, not effective device KV lengths or per-request + ownership. Ownership/incarnation/tickets must be checked separately at publish. + """ + + rank_features: tuple[tuple[int, ...], ...] + capacity: int + actual_rank_rows: tuple[int, ...] + native_rank_rows: tuple[int, ...] + + def __post_init__(self) -> None: + if ( + type(self.rank_features) is not tuple + or len(self.rank_features) != 8 + or type(self.actual_rank_rows) is not tuple + or type(self.native_rank_rows) is not tuple + or len(self.actual_rank_rows) != 8 + or len(self.native_rank_rows) != 8 + or type(self.capacity) is not int + or self.capacity not in _BODY_BUCKETS + ): + raise ValueError("mixed geometry requires eight ordered ranks and a captured capacity") + for rank, features in enumerate(self.rank_features): + if type(features) is not tuple or len(features) != len(_GEOMETRY_FIELDS): + raise ValueError("mixed rank feature dimensions differ") + for value in (*features, self.actual_rank_rows[rank], self.native_rank_rows[rank]): + dspark_planner._require_json_int(value, field="mixed geometry value", minimum=0) + if value >= 2**63: + raise ValueError("mixed geometry exceeds the existing wire integer bound") + ( + context, + _, + _, + _, + history_sum, + history_max, + bootstrap, + real_context, + eligible, + contexts, + generations, + fixed, + ) = features + if ( + real_context > contexts + or bootstrap > generations - eligible + or eligible > generations + or history_max > history_sum + or fixed != (generations - eligible) * 4 + or self.native_rank_rows[rank] != context + fixed + eligible * 4 + or not context + fixed + eligible + <= self.actual_rank_rows[rank] + <= min(self.native_rank_rows[rank], self.capacity) + ): + raise ValueError("mixed geometry does not reserve context and full-K fixed rows") + if ( + not any(row[7] for row in self.rank_features) + or self.actual_rank_rows == self.native_rank_rows + or _paid_bucket(self.native_rank_rows) is None + ): + raise ValueError("geometry is all-decode, noncompact or outside the captured grid") + + def payload(self) -> dict[str, object]: + return { + "rank_features": [dict(zip(_GEOMETRY_FIELDS, row)) for row in self.rank_features], + "capacity": self.capacity, + "actual_rank_rows": list(self.actual_rank_rows), + "native_rank_rows": list(self.native_rank_rows), + } + + +@dataclass(frozen=True) +class MixedEvidenceBinding: + """Exact artifact pins supplied by the optional mixed configuration.""" + + costs_sha256: str + admission_sha256: str + protocol_sha256: str + + def __post_init__(self) -> None: + for field in ("costs_sha256", "admission_sha256", "protocol_sha256"): + dspark_planner._require_sha256(getattr(self, field), field=field) + + +@dataclass(frozen=True) +class MixedCostCell: + geometry: MixedGeometry + native_lower_ms: float + candidate_upper_ms: float + measurement_result_sha256: str + + +@dataclass(frozen=True) +class AdmittedMixedCosts: + identity: MixedRuntimeIdentity + binding: MixedEvidenceBinding + minimum_predicted_gain: float + cells: tuple[MixedCostCell, ...] + + def lookup(self, geometry: MixedGeometry) -> MixedCostCell | None: + """Uncovered/invalid geometry stays native; no nearest-cell selection.""" + if type(geometry) is not MixedGeometry: + return None + return next((cell for cell in self.cells if cell.geometry == geometry), None) + + +@dataclass(frozen=True) +class MixedEvidenceResolution: + costs: AdmittedMixedCosts | None + reason: str + + +def load_mixed_evidence(config) -> MixedEvidenceResolution: + """Read independently pinned deployment artifacts before model execution. + + The live fingerprint is a separately produced artifact, never the table's + identity. Runtime route checks are additional to this byte binding. Invalid + optional evidence leaves mixed verification native, not confidence-off. + """ + if config is None: + return MixedEvidenceResolution(None, "mixed_evidence_absent_native") + try: + binding = MixedEvidenceBinding( + config.costs_sha256, config.admission_sha256, config.protocol_sha256 + ) + + def bounded_file(path: str) -> bytes: + with Path(path).open("rb") as handle: + value = handle.read(_MAX_BYTES + 1) + if not 0 < len(value) <= _MAX_BYTES: + raise ValueError("Mixed artifact exceeds the bounded byte contract") + return value + + live = dspark_planner._validate_exact_fields( + _read_pinned( + bounded_file(config.live_fingerprint_path), + config.live_fingerprint_sha256, + name="mixed live fingerprint", + ), + name="mixed live fingerprint", + fields={"schema", "identity"}, + ) + if live["schema"] != "dspark-mixed-live-fingerprint-v1": + raise ValueError("Unknown independently observed mixed fingerprint") + identity = dspark_planner._validate_exact_fields( + live["identity"], + name="mixed live identity", + fields={ + "source_tree_sha256", + "runtime_snapshot_sha256", + "model_config_sha256", + "gpu_class_sha256", + "topology_sha256", + "selector_identity_sha256", + "physical_k", + "world_size", + "body_buckets", + "execution_mode", + }, + ) + if type(identity["body_buckets"]) is not list: + raise ValueError("Mixed fingerprint body buckets must be a list") + observed = MixedRuntimeIdentity( + **{**identity, "body_buckets": tuple(identity["body_buckets"])} + ) + return resolve_mixed_evidence( + binding, + costs_raw=bounded_file(config.costs_path), + admission_raw=bounded_file(config.admission_path), + live_identity=observed, + ) + except ( + OSError, + TypeError, + ValueError, + KeyError, + AttributeError, + UnicodeError, + OverflowError, + RecursionError, + ) as exc: + return MixedEvidenceResolution( + None, f"mixed_evidence_invalid_native:{type(exc).__name__}:{exc}" + ) + + +def _paid_bucket(rows: tuple[int, ...]) -> int | None: + return next((bucket for bucket in _BODY_BUCKETS if bucket >= max(rows)), None) + + +def _read_pinned(raw: bytes, pin: str, *, name: str) -> dict[str, object]: + if type(raw) is not bytes or not 0 < len(raw) <= _MAX_BYTES: + raise ValueError(f"{name} must be bounded nonempty bytes") + if hashlib.sha256(raw).hexdigest() != pin: + raise ValueError(f"{name} byte digest differs") + + def unique(pairs: list[tuple[str, object]]) -> dict[str, object]: + result = {} + for key, value in pairs: + if key in result: + raise ValueError(f"{name} has duplicate fields") + result[key] = value + return result + + def reject_constant(value: str) -> None: + raise ValueError(f"{name} has nonfinite JSON: {value}") + + value = json.loads(raw, object_pairs_hook=unique, parse_constant=reject_constant) + if not isinstance(value, dict): + raise ValueError(f"{name} must be a JSON object") + return value + + +def _parse_geometry(value: object) -> MixedGeometry: + row = dspark_planner._validate_exact_fields( + value, + name="mixed geometry", + fields={ + "rank_features", + "capacity", + "actual_rank_rows", + "native_rank_rows", + }, + ) + if ( + type(row["rank_features"]) is not list + or type(row["actual_rank_rows"]) is not list + or type(row["native_rank_rows"]) is not list + ): + raise ValueError("mixed geometry arrays differ") + features = tuple( + tuple( + dspark_planner._validate_exact_fields( + rank, name="mixed rank", fields=set(_GEOMETRY_FIELDS) + )[field] + for field in _GEOMETRY_FIELDS + ) + for rank in row["rank_features"] + ) + return MixedGeometry( + features, row["capacity"], tuple(row["actual_rank_rows"]), tuple(row["native_rank_rows"]) + ) + + +def _validate( + binding: MixedEvidenceBinding, + costs_raw: bytes, + admission_raw: bytes, + identity: MixedRuntimeIdentity, +) -> AdmittedMixedCosts: + table = dspark_planner._validate_exact_fields( + _read_pinned(costs_raw, binding.costs_sha256, name="mixed costs"), + name="mixed costs", + fields={ + "schema", + "identity", + "protocol_sha256", + "action_semantics", + "pricing_semantics", + "synthetic", + "provenance", + "minimum_predicted_gain", + "cells", + }, + ) + receipt = dspark_planner._validate_exact_fields( + _read_pinned(admission_raw, binding.admission_sha256, name="mixed admission"), + name="mixed admission", + fields={ + "schema", + "identity", + "costs_sha256", + "protocol_sha256", + "action_semantics", + "pricing_semantics", + "production_admission", + "source_and_measurements_verified", + "accuracy_pass", + "runtime_invariants_pass", + "matched_workload_pass", + "accuracy_result_sha256", + "runtime_invariants_result_sha256", + "matched_workload_result_sha256", + "selector_replay_sha256", + "measurement_results_sha256", + "economics", + }, + ) + if ( + table["schema"] != "dspark-mixed-breakable-costs-v1" + or receipt["schema"] != "dspark-mixed-breakable-admission-v1" + or table["identity"] != identity.payload() + or receipt["identity"] != identity.payload() + or table["protocol_sha256"] != binding.protocol_sha256 + or receipt["protocol_sha256"] != binding.protocol_sha256 + or receipt["costs_sha256"] != binding.costs_sha256 + or any( + payload["action_semantics"] != _ACTION + or payload["pricing_semantics"] != _PRICE_SEMANTICS + for payload in (table, receipt) + ) + or table["synthetic"] is not False + or table["provenance"] != "closed_measured_breakable_mixed" + or any( + receipt[field] is not True + for field in ( + "production_admission", + "source_and_measurements_verified", + "accuracy_pass", + "runtime_invariants_pass", + "matched_workload_pass", + ) + ) + ): + raise ValueError("mixed source, execution, action or production evidence differs") + margin = dspark_planner._require_nonnegative_finite_number( + table["minimum_predicted_gain"], field="mixed gain margin" + ) + if margin < 0.01: + raise ValueError("mixed evidence cannot reduce the incumbent 1% immediate-goodput gate") + for field in ( + "accuracy_result_sha256", + "runtime_invariants_result_sha256", + "matched_workload_result_sha256", + "selector_replay_sha256", + "measurement_results_sha256", + ): + dspark_planner._require_sha256(receipt[field], field=field) + economics = dspark_planner._validate_exact_fields( + receipt["economics"], + name="mixed economics", + fields={ + "policy_steps", + "compact_choices", + "gross_compact_value_ms_lower_bound", + "fixed_confidence_path_overhead_ms_upper_bound", + "safety_margin_ms", + }, + ) + admitted, _ = dspark_planner.evaluate_confidence_workload_admission(**economics) + if not admitted: + raise ValueError("mixed workload economics are not positive") + if type(table["cells"]) is not list or not 0 < len(table["cells"]) <= _MAX_CELLS: + raise ValueError("mixed cells must be a bounded nonempty list") + cells = [] + for item in table["cells"]: + row = dspark_planner._validate_exact_fields( + item, + name="mixed cell", + fields={ + "geometry", + "native_lower_ms", + "candidate_upper_ms", + "measurement_result_sha256", + }, + ) + geometry = _parse_geometry(row["geometry"]) + if any(cell.geometry == geometry for cell in cells): + raise ValueError("duplicate mixed geometry/action cell") + cells.append( + MixedCostCell( + geometry, + dspark_planner._require_positive_finite_number( + row["native_lower_ms"], field="mixed native lower time" + ), + dspark_planner._require_positive_finite_number( + row["candidate_upper_ms"], field="mixed candidate upper time" + ), + dspark_planner._require_sha256( + row["measurement_result_sha256"], field="mixed measurement result" + ), + ) + ) + if ( + dspark_planner._canonical_json_sha256([cell.measurement_result_sha256 for cell in cells]) + != receipt["measurement_results_sha256"] + ): + raise ValueError("mixed admission does not bind all ordered measurement results") + return AdmittedMixedCosts(identity, binding, margin, tuple(cells)) + + +def resolve_mixed_evidence( + binding: MixedEvidenceBinding | None, + *, + costs_raw: bytes | None, + admission_raw: bytes | None, + live_identity: MixedRuntimeIdentity, +) -> MixedEvidenceResolution: + """Resolve once before construction; no decode object/flag is read or mutated. + + Invalid optional mixed evidence fails to native mixed verification. It must + never globally disable an already admitted decode policy, nor be upgraded by + copying a validation-only fingerprint or merely asserting a production flag. + """ + if binding is None: + return MixedEvidenceResolution(None, "mixed_evidence_absent_native") + try: + if ( + type(binding) is not MixedEvidenceBinding + or type(live_identity) is not MixedRuntimeIdentity + ): + raise ValueError("unknown mixed binding/live identity type") + return MixedEvidenceResolution( + _validate(binding, costs_raw, admission_raw, live_identity), "mixed_evidence_bound" + ) + except (TypeError, ValueError, KeyError, UnicodeError, OverflowError, RecursionError) as exc: + return MixedEvidenceResolution( + None, f"mixed_evidence_invalid_native:{type(exc).__name__}:{exc}" + ) diff --git a/tensorrt_llm/_torch/speculative/dspark_planner.py b/tensorrt_llm/_torch/speculative/dspark_planner.py index 812a6932b2db..0d87e590ab4b 100644 --- a/tensorrt_llm/_torch/speculative/dspark_planner.py +++ b/tensorrt_llm/_torch/speculative/dspark_planner.py @@ -175,9 +175,12 @@ class ExactSpsCostTable: max_draft_len: int minimum_predicted_gain: float = 0.01 iteration_drain_guard: ExactSpsDrainGuard | None = None + validate_artifact_freshness: bool = True identity_sha256: str = field(init=False) def __post_init__(self) -> None: + if type(self.validate_artifact_freshness) is not bool: + raise TypeError("validate_artifact_freshness must be a bool") max_draft_len = _require_exact_int(self.max_draft_len, field="max_draft_len", minimum=1) normalized: dict[int, ExactSpsCostRow] = {} for graph_batch_size, table in self.tables.items(): @@ -278,6 +281,9 @@ def __post_init__(self) -> None: for graph_batch_size, table in sorted(normalized.items()) ], } + if not self.validate_artifact_freshness: + # Reuse and strict policies must not share a collective identity. + identity_payload["validate_artifact_freshness"] = False object.__setattr__(self, "identity_sha256", _canonical_json_sha256(identity_payload)) @property @@ -371,36 +377,44 @@ def load_runtime_sps_cost_table( graph_batch_sizes: Sequence[int], max_draft_len: int, live_engine_fingerprint_path: str | Path | None = None, + validate_freshness: bool = True, ) -> tuple[ExactSpsCostTable, dict[str, object]]: - """Load an authenticated schema-v2 exact ``T(G,V)`` table. + """Load structurally checked exact costs, optionally checking freshness. - The source-head pin is intentional: a TRT-LLM version can cover source - changes that alter measured step costs. A changed source identity requires - requalification of the cost table and its workload-admission receipt. + Original artifact metadata is returned unchanged. Reuse does not assert + that recorded timings still describe the active source or GPU node. """ payload = _read_sps_cost_payload(path) schema_version = payload.get("schema_version") if type(schema_version) is not int or schema_version != 2: raise ValueError("Exact SPS cost artifacts require schema_version=2") - if live_engine_fingerprint_path is None: + if validate_freshness and live_engine_fingerprint_path is None: raise ValueError( "Schema-v2 exact SPS artifacts require an independently generated " "live engine fingerprint path" ) - live_engine_fingerprint = _read_sps_cost_payload(live_engine_fingerprint_path) + live_engine_fingerprint = ( + _read_sps_cost_payload(live_engine_fingerprint_path) + if live_engine_fingerprint_path is not None + else None + ) validate_sps_cost_table_payload( payload, graph_batch_sizes=graph_batch_sizes, max_draft_len=max_draft_len, live_engine_fingerprint=live_engine_fingerprint, + validate_freshness=validate_freshness, ) - return _build_exact_sps_cost_table(payload, max_draft_len=max_draft_len), payload + return _build_exact_sps_cost_table( + payload, max_draft_len=max_draft_len, validate_artifact_freshness=validate_freshness + ), payload def _build_exact_sps_cost_table( payload: dict[str, object], *, max_draft_len: int, + validate_artifact_freshness: bool = True, ) -> ExactSpsCostTable: exact_tables = payload["cost_tables"] assert isinstance(exact_tables, dict) @@ -452,6 +466,7 @@ def _build_exact_sps_cost_table( payload["minimum_predicted_gain"], field="minimum_predicted_gain" ), iteration_drain_guard=iteration_drain_guard, + validate_artifact_freshness=validate_artifact_freshness, ) @@ -698,6 +713,7 @@ def validate_sps_cost_table_payload( graph_batch_sizes: Sequence[int], max_draft_len: int, live_engine_fingerprint: dict[str, object] | None = None, + validate_freshness: bool = True, ) -> dict[str, object]: """Validate exact schema-v2 coverage and engine provenance. @@ -711,8 +727,9 @@ def validate_sps_cost_table_payload( graph_batch_sizes: Rank-local CUDA-graph row counts captured by the active engine. max_draft_len: Physical full draft-block length K. - live_engine_fingerprint: Independently supplied runtime fingerprint. - Validation refuses to use an artifact without this comparison. + live_engine_fingerprint: Optional independently supplied runtime identity. + validate_freshness: Require an exact live-identity comparison in addition + to mandatory schema, K, graph coverage and measurement provenance. Returns: The artifact's authenticated engine fingerprint. @@ -733,30 +750,37 @@ def validate_sps_cost_table_payload( assert isinstance(exact_tables, dict) assert isinstance(fingerprint, dict) - if live_engine_fingerprint is None: + if type(validate_freshness) is not bool: + raise TypeError("validate_freshness must be a bool") + if validate_freshness and live_engine_fingerprint is None: raise ValueError( "Schema-v2 SPS validation requires an independently supplied live engine fingerprint" ) - if live_engine_fingerprint is fingerprint: + if validate_freshness and live_engine_fingerprint is fingerprint: raise ValueError( "Live engine fingerprint must be supplied independently of the SPS artifact" ) - live_engine_fingerprint = _validate_exact_fields( - live_engine_fingerprint, - name="live engine fingerprint", - fields=_V2_FINGERPRINT_FIELDS, - ) - _validate_live_fingerprint_values(live_engine_fingerprint) + if live_engine_fingerprint is not None: + live_engine_fingerprint = _validate_exact_fields( + live_engine_fingerprint, + name="live engine fingerprint", + fields=_V2_FINGERPRINT_FIELDS, + ) + _validate_live_fingerprint_values(live_engine_fingerprint) if payload["engine_fingerprint_sha256"] != _canonical_json_sha256(fingerprint): raise ValueError("SPS engine_fingerprint SHA256 does not match payload") - mismatches = sorted( - key for key in _V2_FINGERPRINT_FIELDS if fingerprint[key] != live_engine_fingerprint[key] - ) - if mismatches: - raise ValueError( - "SPS engine_fingerprint does not match active runtime: " + ", ".join(mismatches) + if validate_freshness: + assert live_engine_fingerprint is not None + mismatches = sorted( + key + for key in _V2_FINGERPRINT_FIELDS + if fingerprint[key] != live_engine_fingerprint[key] ) + if mismatches: + raise ValueError( + "SPS engine_fingerprint does not match active runtime: " + ", ".join(mismatches) + ) configured_graph_batch_sizes = { _require_json_int(value, field="runtime graph batch size", minimum=1) @@ -1130,8 +1154,8 @@ def select_exact_sps_candidate( The caller supplies globally agreed expected yields after the phase-2 allgather. This pure layer only compares exact measured cells; it cannot - create a ragged budget or synchronize ranks. Compact serving fails closed - until the exact table carries matched iteration/drain metadata. Every + create a ragged budget or synchronize ranks. Strict validation retains + native verification without matched iteration/drain metadata. Every measured tier that clears the aggregate immediate-goodput guard is reranked by its guarded group-E2E value. Drain loss is the maximum expected loss per real request on any active attention-DP rank, derived from the same @@ -1199,10 +1223,10 @@ def select_exact_sps_candidate( return 0 drain_guard = cost_table.iteration_drain_guard if drain_guard is None: - # The matched mean-output statistic is workload-specific and cannot be - # inferred from T(G,V). An older exact artifact therefore remains a - # valid native-static control but is not allowed to enable compact V. - return 0 + if cost_table.validate_artifact_freshness: + return 0 + # Without optional drain evidence, reuse prices only immediate goodput. + return int(max(goodput_budgets, key=lambda budget: (compact_scores[budget], budget))) native_step_ms = cost_table.step_time(0, graph_batch_size) tail_step_ms = cost_table.step_time(0, drain_guard.tail_graph_batch_size) group_values_ms = {} diff --git a/tensorrt_llm/_torch/speculative/dspark_schedule.py b/tensorrt_llm/_torch/speculative/dspark_schedule.py index e9bd692b6259..2e49e64791b6 100644 --- a/tensorrt_llm/_torch/speculative/dspark_schedule.py +++ b/tensorrt_llm/_torch/speculative/dspark_schedule.py @@ -27,12 +27,16 @@ from dataclasses import dataclass import torch +import triton +import triton.language as tl __all__ = [ "DSparkScheduleConfig", + "DSparkFusedScheduleError", "HOST_POLICY_WINDOWS_SNAPSHOT_OUTPUT", "NATIVE_UNIFORM_VERIFY_OUTPUT", "compute_survival", + "schedule_verify_lens_topk_fused_fill", "schedule_verify_lens_topk", ] @@ -50,6 +54,15 @@ NATIVE_UNIFORM_VERIFY_OUTPUT = "native_uniform_verify" +class DSparkFusedScheduleError(RuntimeError): + """A synchronous optional fused-scheduler dispatch failure. + + Input and graph-contract errors remain ordinary ``ValueError`` exceptions; + callers may safely disable one fused G only for this wrapped launch error + and rerun the established tensor implementation. + """ + + @dataclass(frozen=True) class DSparkScheduleConfig: """Bounds for the per-request verify length. @@ -179,3 +192,269 @@ def schedule_verify_lens_topk( granted = selected.view(bs, schedulable).sum(dim=1).to(torch.int32) verify_lens += granted return torch.clamp(verify_lens, min=floor, max=cfg.resolved_max_verify_len) + + +@triton.jit(do_not_specialize=["num_real", "budget"]) +def _schedule_topk_rank_kernel( + survival_ptr, + output_ptr, + stride_row, + stride_position, + num_real, + budget, + survival_eps, + FLOOR: tl.constexpr, + COLS: tl.constexpr, + NUM_CANDIDATES: tl.constexpr, + BLOCK_COLS: tl.constexpr, + BLOCK_CP: tl.constexpr, +): + request = tl.program_id(0) + position = tl.arange(0, BLOCK_COLS) + candidate_mask = position < COLS + survival = tl.load( + survival_ptr + request * stride_row + (position + FLOOR).to(tl.int64) * stride_position, + mask=candidate_mask, + other=0.0, + ).to(tl.float32) + valid = candidate_mask & (survival >= survival_eps) + masked_survival = tl.where(valid, survival, float("-inf")) + + rank = tl.zeros([BLOCK_COLS], dtype=tl.int32) + for peer_start in range(0, NUM_CANDIDATES, BLOCK_CP): + peer = peer_start + tl.arange(0, BLOCK_CP) + peer_mask = peer < NUM_CANDIDATES + peer_request = peer // COLS + peer_position = peer % COLS + peer_survival = tl.load( + survival_ptr + + peer_request.to(tl.int64) * stride_row + + (peer_position + FLOOR).to(tl.int64) * stride_position, + mask=peer_mask, + other=0.0, + ).to(tl.float32) + peer_valid = peer_mask & (peer_request < num_real) & (peer_survival >= survival_eps) + peer_masked_survival = tl.where(peer_valid, peer_survival, float("-inf")) + + higher = peer_masked_survival[None, :] > masked_survival[:, None] + equal = peer_masked_survival[None, :] == masked_survival[:, None] + earlier = (peer_position[None, :] < position[:, None]) | ( + (peer_position[None, :] == position[:, None]) & (peer_request[None, :] < request) + ) + before = (higher | (equal & earlier)) & peer_valid[None, :] + rank += tl.sum(before.to(tl.int32), axis=1) + + selected = valid & (rank < budget) + tl.store(output_ptr + request, tl.sum(selected.to(tl.int32), axis=0)) + + +@triton.jit(do_not_specialize=["num_real", "budget", "pad_len", "graph_num_tokens"]) +def _schedule_topk_fill_kernel( + output_ptr, + num_rows, + num_real, + budget, + pad_len, + graph_num_tokens, + TOKEN_FLOOR: tl.constexpr, + MAX_TOKEN_LEN: tl.constexpr, + HAS_CANDIDATES: tl.constexpr, + BLOCK: tl.constexpr, +): + rows = tl.arange(0, BLOCK) + row_mask = rows < num_rows + is_real = row_mask & (rows < num_real) + selected_extra = tl.load( + output_ptr + rows, + mask=is_real & (budget > 0) & HAS_CANDIDATES, + other=0, + ).to(tl.int32) + token_len = tl.where(is_real, TOKEN_FLOOR + selected_extra, pad_len) + + # Reproduce fill_bucket_device's exact-pad real-row phase. The explicit + # policy budget controls only the confidence-ranked allocation. Any + # already-paid graph remainder is filled round-robin in request-index + # order; fusion must not promote it into another confidence budget. + spare = graph_num_tokens - tl.sum(tl.where(row_mask, token_len, 0), axis=0) + for _ in range(MAX_TOKEN_LEN - 1): + headroom = is_real & (token_len < MAX_TOKEN_LEN) + in_cycle = tl.cumsum(headroom.to(tl.int32), axis=0) <= spare + grant = headroom & in_cycle + token_len += grant.to(tl.int32) + spare -= tl.sum(grant.to(tl.int32), axis=0) + tl.store(output_ptr + rows, token_len, mask=row_mask) + + +def _schedule_verify_lens_topk_fused_fill_triton( + *, + survival: torch.Tensor, + budget: int, + num_real: int, + pad_len: int, + graph_num_tokens: int, + cfg: DSparkScheduleConfig, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Triton implementation of policy top-k plus exact round-robin fill.""" + num_rows = survival.shape[0] + if num_rows == 0: + return torch.empty(0, dtype=torch.int32, device=survival.device) if out is None else out[:0] + if num_rows > 256: + raise ValueError(f"fused top-k scheduler supports at most 256 rows, got {num_rows}") + schedulable = int(cfg.schedulable_per_request) + output = ( + torch.empty(num_rows, dtype=torch.int32, device=survival.device) + if out is None + else out[:num_rows] + ) + + # Specialize rank work by the captured graph row count, never by the + # step-varying real-row count. This keeps one compiled kernel per all-G + # graph key; peer rows past num_real are masked inside the kernel. + num_candidates = num_rows * schedulable + ranked = budget > 0 and num_real > 0 and num_candidates > 0 + if ranked: + block_cols = triton.next_power_of_2(schedulable) + block_cp = 128 + try: + _schedule_topk_rank_kernel[(num_real,)]( + survival, + output, + survival.stride(0), + survival.stride(1), + num_real, + budget, + float(cfg.survival_eps), + FLOOR=int(cfg.min_verify_len), + COLS=schedulable, + NUM_CANDIDATES=num_candidates, + BLOCK_COLS=block_cols, + BLOCK_CP=block_cp, + ) + except Exception as exc: + raise DSparkFusedScheduleError("DSpark fused top-k rank launch failed") from exc + finalize_block = triton.next_power_of_2(num_rows) + try: + _schedule_topk_fill_kernel[(1,)]( + output, + num_rows, + num_real, + budget, + pad_len, + graph_num_tokens, + TOKEN_FLOOR=int(cfg.min_verify_len) + 1, + MAX_TOKEN_LEN=int(cfg.resolved_max_verify_len) + 1, + HAS_CANDIDATES=schedulable > 0, + BLOCK=finalize_block, + ) + except Exception as exc: + raise DSparkFusedScheduleError("DSpark fused round-robin fill launch failed") from exc + return output + + +def schedule_verify_lens_topk_fused_fill( + *, + survival: torch.Tensor, + budget: int, + num_real: int, + pad_len: int, + cfg: DSparkScheduleConfig, + graph_num_tokens: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Fuse the established top-k schedule and exact round-robin bucket fill. + + The explicit ``budget`` retains exactly the semantics of + :func:`schedule_verify_lens_topk`. Pad rows receive exactly ``pad_len`` + tokens; any graph-token remainder is granted to real rows in index-order + round-robin cycles, exactly as ``fill_bucket_device(..., pad_fill=...)``. + CUDA uses two deterministic Triton launches. CPU executes the established + tensor operations as the functional oracle. When ``out`` is supplied, the + first ``num_rows`` int32 elements are overwritten and returned; CUDA then + performs no output allocation. + """ + if survival.dim() != 2: + raise ValueError(f"survival must be [bs, K], got shape {tuple(survival.shape)}") + num_rows, num_positions = survival.shape + if out is not None: + if out.device != survival.device: + raise ValueError("out must share survival.device") + if out.dtype != torch.int32: + raise TypeError(f"out must have dtype torch.int32, got {out.dtype}") + if out.dim() != 1 or out.numel() < num_rows: + raise ValueError(f"out must be a 1-D tensor with at least {num_rows} elements") + if not out.is_contiguous(): + raise ValueError("out must be contiguous") + if num_positions != cfg.block_size: + raise ValueError( + f"survival has {num_positions} positions but block_size is {cfg.block_size}" + ) + budget = int(budget) + num_real = int(num_real) + pad_len = int(pad_len) + if cfg.survival_eps <= 0.0: + raise ValueError( + "fused top-k scheduling requires survival_eps > 0 so zeroed pad " + "rows cannot compete for the legacy policy budget" + ) + if not 0 <= num_real <= num_rows: + raise ValueError(f"num_real must be in [0, {num_rows}], got {num_real}") + max_token_len = int(cfg.resolved_max_verify_len) + 1 + if not 1 <= pad_len <= max_token_len: + raise ValueError(f"pad_len must be in [1, {max_token_len}], got {pad_len}") + + graph_num_tokens = int(graph_num_tokens) + token_floor = int(cfg.min_verify_len) + 1 + pad_tokens = (num_rows - num_real) * pad_len + minimum_tokens = num_real * token_floor + pad_tokens + maximum_tokens = num_real * max_token_len + pad_tokens + scheduled_capacity = min(max(budget, 0), num_real * int(cfg.schedulable_per_request)) + if not minimum_tokens + scheduled_capacity <= graph_num_tokens <= maximum_tokens: + raise ValueError( + "fused schedule cannot realize the selected CUDA graph without " + "changing policy semantics: " + f"expected [{minimum_tokens + scheduled_capacity}, {maximum_tokens}] " + f"tokens, graph requires {graph_num_tokens}" + ) + + if survival.is_cuda: + return _schedule_verify_lens_topk_fused_fill_triton( + survival=survival, + budget=budget, + num_real=num_real, + pad_len=pad_len, + graph_num_tokens=graph_num_tokens, + cfg=cfg, + out=out, + ) + + rows = torch.arange(num_rows, device=survival.device) + real_survival = torch.where( + (rows < num_real).unsqueeze(1), + survival, + torch.full_like(survival, -1.0), + ) + scheduled = schedule_verify_lens_topk( + survival=real_survival, + budget=budget, + cfg=cfg, + ) + token_lens = torch.where( + rows < num_real, + scheduled + 1, + torch.full_like(scheduled, pad_len), + ) + spare = graph_num_tokens - int(token_lens.sum()) + for _ in range(max_token_len - 1): + if spare <= 0: + break + headroom = (rows < num_real) & (token_lens < max_token_len) + grant = headroom & (torch.cumsum(headroom.to(torch.int64), dim=0) <= spare) + token_lens += grant.to(torch.int32) + spare -= int(grant.sum()) + if spare != 0: + raise ValueError(f"fused schedule left {spare} graph tokens after round-robin fill") + if out is None: + return token_lens + out[:num_rows].copy_(token_lens) + return out[:num_rows] diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py index 69d5b707592d..f955fec01dca 100644 --- a/tensorrt_llm/_torch/speculative/dspark_verify.py +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -124,6 +124,8 @@ def __init__( self._prev_buffer: torch.Tensor | None = None self._prev_event: torch.cuda.Event | None = None self._prev_valid = False + self._host_confidence_stamp: torch.Tensor | None = None + self._prev_confidence_stamp: torch.Tensor | None = None # Count every path that declines to trim; degradation here is silent. self.stats = { "fallback_no_snapshot": 0, @@ -304,7 +306,11 @@ def allocate_exact_sps_candidate( def max_verify_len(self) -> int: return self.cfg.resolved_max_verify_len - def stage_confidence(self, confidence_logits: torch.Tensor) -> None: + def stage_confidence( + self, + confidence_logits: torch.Tensor, + confidence_stamp: torch.Tensor | None = None, + ) -> None: """Start a non-blocking device->host copy of the confidence buffer. ``confidence_logits`` is the worker's whole slot-indexed buffer, staged @@ -317,12 +323,24 @@ def stage_confidence(self, confidence_logits: torch.Tensor) -> None: """ if confidence_logits is None or confidence_logits.shape[0] == 0: return + # Validate before rotating any buffer; a rejected stamp must leave the + # original decode snapshot/event ownership unchanged. + if confidence_stamp is not None and ( + confidence_stamp.dtype != torch.int32 + or confidence_stamp.ndim != 1 + or confidence_stamp.shape[0] != confidence_logits.shape[0] + ): + raise ValueError("Confidence stamps must be int32, one per slot") if self.device_windows: # Rotate before overwriting: what was current becomes the older # (lag-2) snapshot the budget argmax reads; the old older buffer's # storage is reused for this staging. Pure reference swaps. self._host_buffer, self._prev_buffer = (self._prev_buffer, self._host_buffer) self._copy_event, self._prev_event = (self._prev_event, self._copy_event) + self._host_confidence_stamp, self._prev_confidence_stamp = ( + self._prev_confidence_stamp, + self._host_confidence_stamp, + ) self._prev_valid = self._snapshot_valid self._snapshot_valid = False if self._host_buffer is None or self._host_buffer.shape != confidence_logits.shape: @@ -334,6 +352,20 @@ def stage_confidence(self, confidence_logits: torch.Tensor) -> None: ) self._copy_event = torch.cuda.Event() self._host_buffer.copy_(confidence_logits, non_blocking=True) + if confidence_stamp is not None: + if ( + self._host_confidence_stamp is None + or self._host_confidence_stamp.shape != confidence_stamp.shape + ): + self._host_confidence_stamp = torch.empty( + confidence_stamp.shape, + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), + ) + self._host_confidence_stamp.copy_(confidence_stamp, non_blocking=True) + else: + self._host_confidence_stamp = None self._copy_event.record() self._snapshot_valid = True diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 724b2aa15e36..df3a89376dfc 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -1593,6 +1593,52 @@ def invalidate_sampling_params_cache(self) -> None: self._sampling_params_signature[0] = None self._sampling_params_signature[1] = None + def remap_expanded_sampling_params(self, req_idx: torch.Tensor, + num_tokens: int) -> None: + """Repack per-logits-row sampling state after device window selection. + + Device window selection can redistribute a fixed token budget among + requests immediately before replay. The expanded sampling buffers must + follow the resulting token owners. + """ + if num_tokens <= 0: + return + owners = req_idx[:num_tokens].to(torch.long) + pairs = ( + (self.request_temperatures, self.temperatures), + (self.request_top_ks, self.top_ks), + (self.request_top_ps, self.top_ps), + (self.request_min_ps, self.min_ps), + (self.request_seeds, self.seeds), + (self.request_offsets, self.offsets), + ) + for request_values, expanded_values in pairs: + if request_values is None or expanded_values is None: + continue + if expanded_values.numel() < num_tokens: + raise RuntimeError( + "device-selected DSpark layout exceeds the captured " + "sampling-parameter buffer") + torch.index_select(request_values, + 0, + owners, + out=expanded_values[:num_tokens]) + + if self.request_offsets is not None and self.offsets is not None: + if self.qo_indptr is None: + raise RuntimeError( + "device-selected DSpark sampling requires ragged row starts" + ) + row_slots = torch.arange(num_tokens, + device=owners.device, + dtype=torch.int64) + row_slots -= self.qo_indptr.index_select(0, owners) + self.offsets[:num_tokens].add_(row_slots * _RNG_SLOT_SPAN) + + # The expanded cache now follows the device-selected layout. Force the + # next host-prepared iteration to restore its own row expansion. + self._sampling_params_signature[1] = None + class AuxiliarySpeculativeStateHandler(Protocol): """Model-side state that verification writes once per draft token. diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 2626af1ebf23..1c1a1099fd80 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3461,6 +3461,32 @@ def spec_dec_mode(self): return TorchSpeculativeDecodingMode.DFLASH +class DSparkMixedEvidenceConfig(StrictBaseModel): + """Optional measured BREAKABLE mixed evidence, independent of decode costs. + + All SHA256 pins come from the authenticated deployment package. The live + fingerprint must be independently observed for the current runtime; copying + the cost table's identity does not qualify it. Uncovered or invalid mixed + evidence keeps full-width native verification without disabling admitted + decode confidence scheduling. + """ + + costs_path: str = Field( + description="Measured mixed geometry/action cost JSON.") + costs_sha256: str = Field(pattern=r"^[0-9a-f]{64}$", + description="Cost-file SHA256.") + admission_path: str = Field( + description="Qualified mixed admission receipt JSON.") + admission_sha256: str = Field(pattern=r"^[0-9a-f]{64}$", + description="Receipt SHA256.") + protocol_sha256: str = Field(pattern=r"^[0-9a-f]{64}$", + description="Measured protocol SHA256.") + live_fingerprint_path: str = Field( + description="Independently generated mixed runtime fingerprint JSON.") + live_fingerprint_sha256: str = Field(pattern=r"^[0-9a-f]{64}$", + description="Live fingerprint SHA256.") + + class DSparkDecodingConfig(DecodingBaseConfig): """Configuration for DSpark speculative decoding. @@ -3470,9 +3496,8 @@ class DSparkDecodingConfig(DecodingBaseConfig): backbone forward, but it additionally refines the per-position draft logits with a lightweight sequential head (a low-rank Markov head, optionally an RNN head). A confidence head can score each draft position for verification - scheduling without truncating the proposed block. This configuration selects - native full-block drafting and does not expose confidence-scheduling or - calibration controls. + scheduling without truncating the proposed block. Confidence scheduling is + opt-in; the default configuration retains native full-block verification. Key features: - Target-dependent: captures hidden states from ``target_layer_ids``. @@ -3514,6 +3539,89 @@ class DSparkDecodingConfig(DecodingBaseConfig): "read from the draft model config (dspark_markov_head_type), " "defaulting to \"vanilla\".") + enable_confidence_scheduling: bool = Field( + default=False, + description= + "Score each drafted position with the draft's confidence head and let " + "the verification scheduler size the batch's verify budget from those " + "scores. The block is always drafted in full; only the number of tokens " + "sent to the target for verification changes, so acceptance (and hence " + "output distribution) is unaffected. Requires a checkpoint with trained " + "confidence-head weights. Gains only appear at high concurrency -- at " + "small batch sizes the step cost is nearly flat in the verified token " + "count, so trimming saves nothing and costs acceptance length.", + status="prototype") + + validate_confidence_artifact_freshness: bool = Field( + default=False, + description= + "Require exact live cost-artifact identity and a positive workload " + "admission receipt before enabling confidence scheduling. By default, " + "explicitly enabled scheduling reuses structurally valid cost and " + "calibration artifacts, with a warning that performance predictions " + "may be stale. This never changes target verification or acceptance.", + status="prototype") + + confidence_sts_path: Optional[str] = Field( + default=None, + description= + "Path to a JSON file with per-position sequential-temperature-scaling " + "(STS) values used to calibrate the confidence head. The scheduler " + "consumes a cumulative product of per-position probabilities, so a " + "per-position calibration error compounds geometrically along the " + "block. If None, the raw sigmoid is used (identity calibration).", + status="prototype") + + confidence_sps_table_path: Optional[str] = Field( + default=None, + description= + "Path to an authenticated schema-v2 JSON table of measured decode step " + "costs keyed by exact rank-local CUDA-graph shape (G, V).", + status="prototype") + + confidence_sps_live_fingerprint_path: Optional[str] = Field( + default=None, + description= + "Path to an independently generated JSON fingerprint of the active " + "runtime. Required when validate_confidence_artifact_freshness=True. " + "Otherwise it is optional observed provenance, kept separate from " + "the cost artifact's original metadata.", + status="prototype") + + confidence_admission_receipt_path: Optional[str] = Field( + default=None, + description= + "Path to an authenticated workload-level confidence admission receipt. " + "Required only when validate_confidence_artifact_freshness=True; it " + "must prove positive aggregate compact value after the fixed " + "confidence-path overhead and safety margin are charged.", + status="prototype") + + confidence_admission_receipt_sha256: Optional[str] = Field( + default=None, + description="Externally authenticated lowercase SHA256 of " + "confidence_admission_receipt_path. This pin must come from the sealed " + "deployment package rather than from the receipt itself.", + status="prototype") + + enable_fused_confidence_scheduler: bool = Field( + default=False, + description= + "Use the optional policy-neutral Triton implementation of confidence " + "top-k allocation plus exact CUDA-graph bucket fill. This changes only " + "the scheduler launch path: budgets, survival thresholds, selected " + "windows, and fallback behavior remain identical. It is independently " + "default-off and requires confidence scheduling.", + status="prototype") + + confidence_mixed_evidence: Optional[DSparkMixedEvidenceConfig] = Field( + default=None, + description="Optional separately admitted physical-K3, eight-rank " + "BREAKABLE mixed scheduling evidence. Does not reuse decode T(G,V) costs. " + "None, unsupported routes, stale ownership, or uncovered mixed geometry " + "retain native mixed verification while admitted decode remains enabled.", + status="prototype") + decoding_type: Literal["DSpark"] = Field(default="DSpark") attention_backend: Literal["AUTO", "VANILLA", "TRTLLM", "CUTEDSL"] = Field( @@ -3529,6 +3637,51 @@ class DSparkDecodingConfig(DecodingBaseConfig): "and needs a cute-dsl MLA decode taking per-token kv_bounds that is not " "upstream yet. Which kernel each name selects: MLADSparkForCausalLM.") + @model_validator(mode="after") + def validate_confidence_scheduling(self): + if not self.enable_confidence_scheduling: + if (self.validate_confidence_artifact_freshness + or self.enable_fused_confidence_scheduler + or self.confidence_sts_path + or self.confidence_sps_table_path + or self.confidence_sps_live_fingerprint_path + or self.confidence_admission_receipt_path + or self.confidence_admission_receipt_sha256 + or self.confidence_mixed_evidence is not None): + raise ValueError( + "validate_confidence_artifact_freshness / " + "enable_fused_confidence_scheduler / confidence_sts_path / " + "confidence_sps_table_path / " + "confidence_sps_live_fingerprint_path / " + "confidence_admission_receipt_path / " + "confidence_admission_receipt_sha256 / confidence_mixed_evidence require " + "enable_confidence_scheduling=True") + return self + + if self.max_draft_len == 1: + # Physical K1 has no smaller verifier window, so confidence + # scheduling can only add control-path work. Resolve it to the + # ordinary feature-off configuration before any runtime component + # observes the config. This preserves the original static K path + # rather than running a confidence-enabled native fallback. + logger.warning( + "DSpark confidence scheduling has no compact tier below the " + "native K1 layout; using the feature-off static K1 path and " + "ignoring confidence-scheduling inputs.") + self.enable_confidence_scheduling = False + self.validate_confidence_artifact_freshness = False + self.enable_fused_confidence_scheduler = False + self.confidence_sts_path = None + self.confidence_sps_table_path = None + self.confidence_sps_live_fingerprint_path = None + self.confidence_admission_receipt_path = None + self.confidence_admission_receipt_sha256 = None + self.confidence_mixed_evidence = None + return self + # Resolve optional freshness admission before runtime components + # snapshot the flag. Exact graph/K coverage is checked by the engine. + return self + @model_validator(mode="after") def set_max_total_draft_tokens(self): self.max_total_draft_tokens = self.max_draft_len @@ -6741,6 +6894,170 @@ def validate_encode_only_torch_compile_config(self) -> 'TorchLlmArgs': "graphs or disable enable_piecewise_cuda_graph.") return self + def _validate_dspark_confidence_environment(self) -> None: + """Reject configurations confidence scheduling cannot honour. + + Every check here guards a path that would otherwise degrade *silently*: + the run comes up, produces plausible text, and either never takes the + ragged path at all or takes it with a kernel that mis-attributes rows. + A hard error at construction is the only signal that survives. + """ + if not self.speculative_config.draft_is_embedded_in_target: + raise ValueError( + "DSpark confidence scheduling supports only the embedded " + "DeepSeek-V4 draft. Standalone DSpark retains native " + "verification; disable confidence scheduling for that path.") + cuda_graph_config = self.cuda_graph_config + if self.pipeline_parallel_size > 1: + raise ValueError( + "speculative_config.enable_confidence_scheduling=True does not yet " + "support pipeline parallelism. The PP executor's pipelined " + "microbatches do not provide the single-iteration host-input " + "lifetime boundary required by ragged verification staging. " + "Set pipeline_parallel_size=1 or disable confidence scheduling." + ) + if self.guided_decoding_backend is not None: + raise ValueError( + "speculative_config.enable_confidence_scheduling=True does not yet " + "support guided decoding. Guided-decoder bitmasks are laid " + "out with one uniform speculative stride; applying them to " + "packed per-request verify windows would shift every request " + "after the first short window. Disable guided decoding or " + "disable confidence scheduling.") + if self.enable_lora or self.lora_config is not None: + raise ValueError( + "speculative_config.enable_confidence_scheduling=True does not yet " + "support LoRA. LoRA metadata expands generation requests " + "with one uniform speculative stride, which cannot represent " + "packed per-request verify windows. Disable LoRA or disable " + "confidence scheduling.") + # `fit_ragged_verify_lens` plans against the padded row count and + # reserves tokens for the CUDA-graph padding rows, so it presumes those + # rows exist. With padding off, `_get_padded_batch` adds none: the batch + # then carries a token total that keys into a graph nobody captured and + # drops to eager on every ragged step -- which costs far more than the + # trimmed tokens save -- while real requests were still shrunk to leave + # room for padding that never arrived. + if cuda_graph_config is None: + raise ValueError( + "speculative_config.enable_confidence_scheduling=True requires " + "cuda_graph_config to be set with enable_padding=True. Ragged " + "verification fits each step's token total onto a captured " + "CUDA-graph bucket; with no graphs there is no bucket to fit " + "and every step would silently run eager.") + if not cuda_graph_config.enable_padding: + raise ValueError( + "speculative_config.enable_confidence_scheduling=True requires " + "cuda_graph_config.enable_padding=True (it defaults to False). " + "The ragged token budget is planned against the padded batch " + "size and reserves tokens for the padding rows; without padding " + "those rows never appear, so the step's token total misses its " + "captured bucket and falls out of graph replay.") + + sparse_attention_config = self.sparse_attention_config + if sparse_attention_config is not None: + # Both recover "which request is this row" from a fixed next_n -- + # GVR indexes its previous-step top-k hint by request and strides it + # by next_n, and the CuTe-DSL kernel takes next_n as a compile-time + # constant and JIT cache key. Neither can express per-request + # windows. The C++ launcher already refuses the combination + # (indexerTopK.cu, "ragged rowKvLens is incompatible with the GVR + # preIdx hint"), but that fires deep in a decode step; say so here. + incompatible = [ + name for name in ("enable_heuristic_topk", "use_cute_dsl_topk") + if getattr(sparse_attention_config, name, False) + ] + if incompatible: + raise ValueError( + "speculative_config.enable_confidence_scheduling=True is " + "incompatible with sparse_attention_config." + f"{' / '.join(incompatible)}=True: those top-k paths " + "reconstruct each row's request and window position from a " + "single batch-wide next_n, which a ragged batch does not " + "have. Disable them, or disable confidence scheduling.") + # Not an error: the DSL paged-MQA-logits path is structurally + # uniform too, but the expanded (one row per query token) path it + # falls back to is numerically equivalent, and it is the default on + # DeepSeek-V4 anyway. Still worth saying out loud, since the knob + # will read as honoured otherwise. + if getattr(sparse_attention_config, "use_cute_dsl_paged_mqa_logits", + False): + logger.warning( + "sparse_attention_config.use_cute_dsl_paged_mqa_logits=True " + "is ignored on ragged-verification steps: the DSL kernel " + "broadcasts one KV length per request across a fixed next_n. " + "Those steps use the expanded DeepGEMM path instead, which " + "is numerically equivalent.") + + @staticmethod + def _resolve_dspark_confidence_workload_admission( + spec_cfg: DSparkDecodingConfig) -> None: + """Resolve confidence activation before engines or workers observe it.""" + from tensorrt_llm._torch.speculative.dspark_planner import \ + load_confidence_workload_admission + + receipt_path = spec_cfg.confidence_admission_receipt_path + receipt_sha256 = spec_cfg.confidence_admission_receipt_sha256 + sps_table_path = spec_cfg.confidence_sps_table_path + live_fingerprint_path = spec_cfg.confidence_sps_live_fingerprint_path + if not spec_cfg.validate_confidence_artifact_freshness: + if not sps_table_path: + raise ValueError( + "DSpark confidence scheduling requires confidence_sps_table_path; " + "a missing measured cost table cannot be timed as enabled.") + logger.warning( + "DSpark confidence scheduling is enabled with reusable artifacts. " + "Cost/calibration predictions and performance economics are not " + "freshness-validated and may be stale; target acceptance is unchanged. " + "Set validate_confidence_artifact_freshness=True to require exact " + "live provenance and positive workload admission.") + return + missing = [ + name for name, value in ( + ("confidence_admission_receipt_path", receipt_path), + ("confidence_admission_receipt_sha256", receipt_sha256), + ("confidence_sps_table_path", sps_table_path), + ("confidence_sps_live_fingerprint_path", live_fingerprint_path), + ) if not value + ] + fallback_reason = ("missing " + ", ".join(missing)) if missing else None + if fallback_reason is None: + assert receipt_path is not None + assert receipt_sha256 is not None + assert sps_table_path is not None + assert live_fingerprint_path is not None + try: + admission = load_confidence_workload_admission( + receipt_path, + expected_receipt_sha256=receipt_sha256, + sps_cost_table_path=sps_table_path, + live_engine_fingerprint_path=live_fingerprint_path, + physical_k=spec_cfg.max_draft_len, + ) + if not admission.admitted: + fallback_reason = ( + "the authenticated workload replay does not have positive " + "net compact value after fixed overhead and safety margin" + ) + except (OSError, TypeError, ValueError) as error: + fallback_reason = str(error) + + if fallback_reason is None: + return + + logger.warning( + f"DSpark confidence scheduling is not deployment-admitted ({fallback_reason}); " + f"using the feature-off static K{spec_cfg.max_draft_len} path.") + spec_cfg.enable_confidence_scheduling = False + spec_cfg.validate_confidence_artifact_freshness = False + spec_cfg.enable_fused_confidence_scheduler = False + spec_cfg.confidence_sts_path = None + spec_cfg.confidence_sps_table_path = None + spec_cfg.confidence_sps_live_fingerprint_path = None + spec_cfg.confidence_admission_receipt_path = None + spec_cfg.confidence_admission_receipt_sha256 = None + spec_cfg.confidence_mixed_evidence = None + @model_validator(mode="after") def normalize_prefill_cuda_graph_config(self) -> 'TorchLlmArgs': """Normalize legacy piecewise CUDA graph options into prefill fields.""" @@ -7211,6 +7528,10 @@ def _dspark_get( "DSpark block_size must equal max_draft_len; got " f"block_size={spec_cfg.block_size} and " f"max_draft_len={spec_cfg.max_draft_len}") + if spec_cfg.enable_confidence_scheduling: + self._resolve_dspark_confidence_workload_admission(spec_cfg) + if spec_cfg.enable_confidence_scheduling: + self._validate_dspark_confidence_environment() if isinstance(self.speculative_config, SADecodingConfig): pool_size = self.speculative_config.global_pool_size diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 9b9e887b5b98..ed49a7d27f32 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1674,6 +1674,16 @@ "kind": "categorical", "path": "speculative_config.eagle3_model_arch" }, + { + "capture_policy": "bool", + "kind": "value", + "path": "speculative_config.enable_confidence_scheduling" + }, + { + "capture_policy": "bool", + "kind": "value", + "path": "speculative_config.enable_fused_confidence_scheduler" + }, { "capture_policy": "bool", "kind": "value", @@ -1833,6 +1843,11 @@ "kind": "value", "path": "speculative_config.use_relaxed_acceptance_for_thinking" }, + { + "capture_policy": "bool", + "kind": "value", + "path": "speculative_config.validate_confidence_artifact_freshness" + }, { "capture_policy": "int", "kind": "value", diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index da00b2aa01bf..b3646f2f7ea6 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -88,6 +88,7 @@ l0_cpu: - unittest/_torch/multimodal - unittest/_torch/ray_orchestrator/single_gpu/test_cache_transceiver_comm.py - unittest/_torch/speculative/hw_agnostic + - unittest/_torch/speculative/test_device_window_sampling_layout.py # l0_h100 selects with -m "not cpu_only", so this entry is what runs it. - unittest/_torch/speculative/test_spec_slot_pool_sizing.py - unittest/_torch/test_hf_quant_config.py @@ -188,6 +189,7 @@ l0_cpu: - unittest/llmapi/test_whisper_suppress_tokens_processor.py - unittest/llmapi/test_kv_cache_dtype_override.py - unittest/llmapi/test_llm_args.py + - unittest/llmapi/test_dspark_workload_admission_config.py - unittest/llmapi/test_mx_args.py - unittest/llmapi/test_custom_tokenizer_aliases.py - unittest/llmapi/test_llm_quant.py diff --git a/tests/unittest/_torch/executor/engine/test_metadata.py b/tests/unittest/_torch/executor/engine/test_metadata.py index 801d01b70f2f..392afda5f521 100644 --- a/tests/unittest/_torch/executor/engine/test_metadata.py +++ b/tests/unittest/_torch/executor/engine/test_metadata.py @@ -95,3 +95,47 @@ def test_build_attention_metadata_forwards_shared_and_cache_inputs() -> None: assert metadata.draft_kv_cache_manager is draft_kv_cache_manager assert metadata.enable_context_mla_with_cached_kv assert metadata.num_heads_per_kv == 4 + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_build_attention_metadata_passes_ragged_opt_in_to_supported_backend(enabled) -> None: + class RaggedMetadata(_AttentionMetadata): + enable_ragged_verification = False + + model_config = SimpleNamespace( + pretrained_config=SimpleNamespace(), + sparse_attention_config=None, + enable_flash_mla=False, + ) + metadata = build_attention_metadata( + model_config, + max_batch_size=4, + max_num_tokens=16, + max_beam_width=1, + attention_backend=SimpleNamespace(Metadata=RaggedMetadata), + attention_runtime_features=AttentionRuntimeFeatures(), + mapping=object(), + cache_indirection=None, + enable_ragged_verification=enabled, + ) + assert metadata.enable_ragged_verification is enabled + + +def test_build_attention_metadata_does_not_pass_ragged_option_to_other_backend() -> None: + model_config = SimpleNamespace( + pretrained_config=SimpleNamespace(), + sparse_attention_config=None, + enable_flash_mla=False, + ) + metadata = build_attention_metadata( + model_config, + max_batch_size=4, + max_num_tokens=16, + max_beam_width=1, + attention_backend=_AttentionBackend, + attention_runtime_features=AttentionRuntimeFeatures(), + mapping=object(), + cache_indirection=None, + enable_ragged_verification=True, + ) + assert not hasattr(metadata, "enable_ragged_verification") diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py index f68350399cf5..404c0ca5fbcd 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py @@ -8,6 +8,7 @@ import tensorrt_llm import tensorrt_llm.bindings +from tensorrt_llm._torch.pyexecutor._util import _get_num_reserved_index_slots from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2 as kv_cache_v2_module from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState, SamplingConfig @@ -16,6 +17,27 @@ CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType +@pytest.mark.parametrize( + ("spec_config", "estimating_kv_cache", "expected"), + [ + (None, False, 1), + (SimpleNamespace(enable_confidence_scheduling=False), False, 1), + (SimpleNamespace(enable_confidence_scheduling=True), False, 2), + (SimpleNamespace(enable_confidence_scheduling=True), True, 0), + ], + ids=[ + "no-speculation", + "static-dspark", + "confidence-dspark", + "confidence-estimator", + ], +) +def test_index_mapper_reservation_tracks_confidence_scheduling( + spec_config: SimpleNamespace | None, estimating_kv_cache: bool, expected: int +) -> None: + assert _get_num_reserved_index_slots(spec_config, estimating_kv_cache) == expected + + def _manager( *, is_draft: bool, diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index ad0cb0482ac4..7985df7c24fc 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -768,6 +768,7 @@ def test_graph_key_forwards_promoted_context_ids(self) -> None: runner.config = SimpleNamespace() runner._resolve_sample_type.return_value = SampleType.FULL runner._get_seq_len_mode.return_value = True + runner._ragged_verify_bucket.return_value = None request = _make_request_stub(7) batch = ScheduledRequests() batch.generation_requests = [request] @@ -796,6 +797,7 @@ def test_graph_key_aggregates_encoder_tokens(self) -> None: runner._resolve_sample_type.return_value = SampleType.FULL runner.max_beam_width = 1 runner._get_seq_len_mode.return_value = False + runner._ragged_verify_bucket.return_value = None context = _make_request_stub(1) context.encoder_output_len = 7 context.py_skip_cross_kv_projection = False @@ -820,6 +822,7 @@ def test_graph_key_rejects_nonuniform_context_query_lengths(self) -> None: runner.config = SimpleNamespace() runner._resolve_sample_type.return_value = SampleType.FULL runner._get_seq_len_mode.return_value = False + runner._ragged_verify_bucket.return_value = None first_context = _make_request_stub(1) first_context.encoder_output_len = 7 first_context.py_skip_cross_kv_projection = False @@ -867,6 +870,7 @@ def test_graph_key_includes_peft_cache_dtype(self) -> None: runner.config = SimpleNamespace() runner._resolve_sample_type.return_value = SampleType.FULL runner._get_seq_len_mode.return_value = False + runner._ragged_verify_bucket.return_value = None request = _make_request_stub(7) batch = ScheduledRequests() batch.generation_requests = [request] @@ -885,6 +889,7 @@ def test_graph_key_includes_peft_cache_dtype(self) -> None: def test_graph_dtype_change_falls_back_to_eager(self) -> None: runner = Mock() runner.enabled = True + runner._dspark_confidence_enabled = False runner.config = SimpleNamespace( enable_attention_dp=False, use_mrope=False, @@ -925,6 +930,7 @@ def test_graph_key_includes_lora_variant(self) -> None: runner.config = SimpleNamespace() runner._resolve_sample_type.return_value = SampleType.FULL runner._get_seq_len_mode.return_value = False + runner._ragged_verify_bucket.return_value = None request = _make_request_stub(7) batch = ScheduledRequests() batch.generation_requests = [request] @@ -967,6 +973,7 @@ def test_lora_graph_variant_selection(self) -> None: def test_graph_lookup_forwards_promoted_context_ids(self) -> None: runner = Mock() runner.enabled = True + runner._dspark_confidence_enabled = False runner.config = SimpleNamespace( enable_attention_dp=False, use_mrope=False, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_cuda_graph_dummy_identity.py b/tests/unittest/_torch/speculative/hw_agnostic/test_cuda_graph_dummy_identity.py new file mode 100644 index 000000000000..2846cb2f50dc --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_cuda_graph_dummy_identity.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( + CUDA_GRAPH_DUMMY_REQUEST_ID, + _padding_dummy_draft_token_capacity, + cuda_graph_dummy_request_id, +) +from tensorrt_llm._torch.speculative import dspark as dspark_module +from tensorrt_llm._torch.speculative.dspark import DSparkSpecMetadata, DSv4DSparkWorker + +pytestmark = pytest.mark.cpu_only + + +@pytest.mark.parametrize("mode", ["Eagle", "MTP"]) +def test_tree_speculation_keeps_native_dummy_identity(mode): + config = SimpleNamespace(decoding_type=mode, max_draft_len=3) + capacity = _padding_dummy_draft_token_capacity(config, 12) + assert capacity == 12 + assert cuda_graph_dummy_request_id(12, variant=0, max_draft_len=capacity) == ( + CUDA_GRAPH_DUMMY_REQUEST_ID - 12 + ) + + +def test_dspark_does_not_relax_the_physical_k_bound(): + config = SimpleNamespace(decoding_type="DSpark", max_draft_len=3) + capacity = _padding_dummy_draft_token_capacity(config, 12) + assert capacity == 3 + with pytest.raises(ValueError, match="invalid CUDA padding dummy identity"): + cuda_graph_dummy_request_id(12, variant=0, max_draft_len=capacity) + + +def test_secondary_dspark_dummy_ids_are_disjoint(): + native = {cuda_graph_dummy_request_id(k, variant=0, max_draft_len=3) for k in range(4)} + secondary = {cuda_graph_dummy_request_id(k, variant=1, max_draft_len=3) for k in range(4)} + assert native.isdisjoint(secondary) + + +def test_target_only_dummy_identity_is_unchanged(): + capacity = _padding_dummy_draft_token_capacity(None, 0) + assert cuda_graph_dummy_request_id(0, variant=0, max_draft_len=capacity) == ( + CUDA_GRAPH_DUMMY_REQUEST_ID + ) + + +@pytest.mark.parametrize("physical_k", [1, 3, 8]) +def test_padding_families_never_consume_real_worker_slots(monkeypatch, physical_k): + def cpu_allocator(original): + def allocate(*args, **kwargs): + kwargs["device"] = "cpu" + return original(*args, **kwargs) + + return allocate + + for name in ("empty", "zeros", "full"): + monkeypatch.setattr(torch, name, cpu_allocator(getattr(torch, name))) + monkeypatch.setattr(dspark_module, "prefer_pinned", lambda: False) + worker = object.__new__(DSv4DSparkWorker) + worker.spec_config = SimpleNamespace(max_draft_len=physical_k) + worker._win_inited = False + worker._attention_warmup_attempted = True + worker.return_confidence = False + worker._confidence_logits = None + worker._mixed_confidence_tracking = False + worker._confidence_incarnations = {} + draft = SimpleNamespace( + block_size=physical_k, + num_stages=1, + _attn_params={"window_size": 4, "head_dim": 2}, + ) + worker._lazy_init(draft, SimpleNamespace(max_num_requests=4)) + metadata = SimpleNamespace( + _dspark_worker=worker, + batch_indices_cuda=torch.empty(4, dtype=torch.int32), + ) + + for variant in (0, 1): + for runtime_k in range(physical_k + 1): + dummy_id = cuda_graph_dummy_request_id( + runtime_k, variant=variant, max_draft_len=physical_k + ) + metadata.request_ids = [1000, 0, dummy_id] + metadata.num_generations = 3 + DSparkSpecMetadata.prepare(metadata) + assert worker._req_to_slot == {1000: 0} + assert list(worker._free_slots) == [1, 2, 3] + assert worker._batch_to_slot[:3].tolist() == [0, 4, 4] + + metadata.request_ids = [0, dummy_id] + metadata.num_generations = 2 + DSparkSpecMetadata.prepare(metadata) + assert worker._req_to_slot == {} + assert sorted(worker._free_slots) == [0, 1, 2, 3] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_admission_resolution.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_admission_resolution.py new file mode 100644 index 000000000000..e35300747444 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_admission_resolution.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Configuration-level native fallback, before any worker snapshots its flags.""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.speculative import dspark_planner +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig, TorchLlmArgs + +pytestmark = pytest.mark.cpu_only + +_ARTIFACT_FIELDS = ( + "confidence_sts_path", + "confidence_sps_table_path", + "confidence_sps_live_fingerprint_path", + "confidence_admission_receipt_path", + "confidence_admission_receipt_sha256", +) + + +def _config(physical_k=3, **kwargs): + return DSparkDecodingConfig( + max_draft_len=physical_k, + enable_confidence_scheduling=True, + validate_confidence_artifact_freshness=True, + enable_fused_confidence_scheduler=True, + **kwargs, + ) + + +@pytest.mark.parametrize("outcome", ["missing", "negative", "stale", "unreadable", "positive"]) +@pytest.mark.parametrize("physical_k", [2, 3, 4, 5, 8]) +def test_admission_resolves_before_worker_construction(monkeypatch, physical_k, outcome): + artifacts = {name: "/sealed/artifact.json" for name in _ARTIFACT_FIELDS} + artifacts["confidence_admission_receipt_sha256"] = "a" * 64 + config = _config(physical_k, **({} if outcome == "missing" else artifacts)) + calls = [] + + def load(*args, **kwargs): + calls.append((args, kwargs)) + if outcome == "stale": + raise ValueError("source identity mismatch") + if outcome == "unreadable": + raise OSError("artifact unavailable") + return SimpleNamespace(admitted=outcome == "positive") + + monkeypatch.setattr(dspark_planner, "load_confidence_workload_admission", load) + TorchLlmArgs._resolve_dspark_confidence_workload_admission(config) + assert config.max_draft_len == physical_k + assert config.tokens_per_gen_step == physical_k + 1 + assert config.enable_confidence_scheduling is (outcome == "positive") + assert config.enable_fused_confidence_scheduler is (outcome == "positive") + # Experiment classification must use the resolved snapshot, not the input + # enable flag. A successful fallback is not an enabled-native experiment. + resolved = config.model_dump() + assert resolved["enable_confidence_scheduling"] is (outcome == "positive") + assert resolved["enable_fused_confidence_scheduler"] is (outcome == "positive") + for name in _ARTIFACT_FIELDS: + assert getattr(config, name) == (artifacts[name] if outcome == "positive" else None) + assert resolved[name] == (artifacts[name] if outcome == "positive" else None) + if outcome == "missing": + assert not calls + else: + assert calls[0][1]["physical_k"] == physical_k + assert calls[0][1]["expected_receipt_sha256"] == "a" * 64 + + +def test_k1_never_needs_admission_or_confidence_artifacts(): + config = _config(1) + assert config.enable_confidence_scheduling is False + assert config.enable_fused_confidence_scheduler is False + assert config.tokens_per_gen_step == 2 + assert all(getattr(config, name) is None for name in _ARTIFACT_FIELDS) + + +def test_explicit_enable_defaults_to_reuse_without_inventing_admission(monkeypatch): + config = DSparkDecodingConfig( + max_draft_len=3, + enable_confidence_scheduling=True, + confidence_sps_table_path="/original/sps.json", + confidence_sts_path="/original/sts.json", + ) + + def reject_load(*args, **kwargs): + raise AssertionError("Reuse must not fabricate or require an admission receipt") + + monkeypatch.setattr(dspark_planner, "load_confidence_workload_admission", reject_load) + TorchLlmArgs._resolve_dspark_confidence_workload_admission(config) + assert config.enable_confidence_scheduling is True + assert config.validate_confidence_artifact_freshness is False + assert config.confidence_admission_receipt_path is None + assert config.confidence_sps_table_path == "/original/sps.json" + restored = DSparkDecodingConfig.model_validate_json(config.model_dump_json()) + assert restored.enable_confidence_scheduling is True + assert restored.validate_confidence_artifact_freshness is False + assert restored.confidence_sts_path == "/original/sts.json" + + +def test_reuse_refuses_to_claim_enabled_without_a_cost_table(): + config = DSparkDecodingConfig(max_draft_len=3, enable_confidence_scheduling=True) + with pytest.raises(ValueError, match="requires confidence_sps_table_path"): + TorchLlmArgs._resolve_dspark_confidence_workload_admission(config) + + +def test_default_off_and_strict_mode_round_trip(): + config = DSparkDecodingConfig(max_draft_len=3) + assert config.enable_confidence_scheduling is False + assert config.validate_confidence_artifact_freshness is False + strict = _config(3) + restored = DSparkDecodingConfig.model_validate_json(strict.model_dump_json()) + assert restored.validate_confidence_artifact_freshness is True + with pytest.raises(ValueError, match="enable_confidence_scheduling=True"): + DSparkDecodingConfig(max_draft_len=3, validate_confidence_artifact_freshness=True) + + +@pytest.mark.parametrize("field_name", _ARTIFACT_FIELDS) +def test_feature_off_rejects_confidence_only_inputs(field_name): + with pytest.raises(ValueError, match="enable_confidence_scheduling=True"): + DSparkDecodingConfig(max_draft_len=3, **{field_name: "unconsumed"}) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_artifact_reuse.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_artifact_reuse.py new file mode 100644 index 000000000000..65d94ac666d0 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_artifact_reuse.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Cost reuse changes predictions/activation, never target acceptance.""" + +import copy +import hashlib +import json +from pathlib import Path + +import pytest +from pydantic import BaseModel + +from tensorrt_llm._torch.speculative.dspark_planner import ( + ExactSpsCostRow, + ExactSpsCostTable, + ExactSpsDrainGuard, + load_runtime_sps_cost_table, + select_exact_sps_candidate, +) +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig +from tensorrt_llm.usage.llmapi_config import collect_llm_api_config_payloads, manifest_rows + +pytestmark = pytest.mark.cpu_only + + +def _payload(): + fingerprint = { + "gpu": "B300", + "gpu_count": 8, + "gpu_snapshot_sha256": "a" * 64, + "global_graph_batch_sizes": [512], + "max_draft_len": 3, + "rank_local_graph_batch_sizes": [64], + "runtime_snapshot": "original-runtime", + "source_diff_sha256": "b" * 64, + "source_head": "original-source", + "topology": "DEP8", + } + return { + "schema_version": 2, + "minimum_predicted_gain": 0.01, + "engine_fingerprint": fingerprint, + "engine_fingerprint_sha256": hashlib.sha256( + json.dumps(fingerprint, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest(), + "cost_tables": {"64": {"token_counts": [0, 128], "step_time_ms": [10.0, 8.0]}}, + "measurements": [ + { + "rank_local_graph_batch_size": 64, + "rank_local_verifier_budget": budget, + "step_time_ms": step, + "source_result_sha256": "c" * 64, + } + for budget, step in ((0, 10.0), (128, 8.0)) + ], + } + + +def _load(tmp_path, payload, *, strict=False, live=None): + path = tmp_path / "sps.json" + path.write_text(json.dumps(payload), encoding="utf-8") + live_path = None + if live is not None: + live_path = tmp_path / "live.json" + live_path.write_text(json.dumps(live), encoding="utf-8") + return load_runtime_sps_cost_table( + path, + graph_batch_sizes=[64], + max_draft_len=3, + live_engine_fingerprint_path=live_path, + validate_freshness=strict, + ) + + +def _select(table, *, yield_value=9.0): + return select_exact_sps_candidate( + graph_batch_size=64, + native_expected_yield=10.0, + compact_expected_yields={128: yield_value}, + compact_max_yield_losses_per_request={128: 10.0 - yield_value}, + cost_table=table, + ) + + +def test_reuse_without_live_receipt_or_drain_can_genuinely_select_compact(tmp_path): + payload = _payload() + before = copy.deepcopy(payload) + table, metadata = _load(tmp_path, payload) + assert metadata == before + assert table.validate_artifact_freshness is False + assert table.iteration_drain_guard is None + assert _select(table) == 128 + assert _select(table, yield_value=8.0) == 0 + assert table.minimum_predicted_gain == 0.01 + + +@pytest.mark.parametrize( + "field", ["source_head", "runtime_snapshot", "source_diff_sha256", "gpu_snapshot_sha256"] +) +def test_reuse_retains_original_metadata_but_strict_rejects_staleness(tmp_path, field): + payload = _payload() + live = copy.deepcopy(payload["engine_fingerprint"]) + live[field] = "d" * 64 + table, metadata = _load(tmp_path, payload, live=live) + assert _select(table) == 128 + assert metadata["engine_fingerprint"] == payload["engine_fingerprint"] + assert metadata["engine_fingerprint"][field] != live[field] + with pytest.raises(ValueError, match="does not match active runtime"): + _load(tmp_path, payload, strict=True, live=live) + + +def test_strict_requires_independent_live_and_retains_absent_drain_refusal(tmp_path): + payload = _payload() + with pytest.raises(ValueError, match="independently generated"): + _load(tmp_path, payload, strict=True) + table, _ = _load( + tmp_path, payload, strict=True, live=copy.deepcopy(payload["engine_fingerprint"]) + ) + assert _select(table) == 0 + + +@pytest.mark.parametrize( + "fault", + [ + "nonfinite", + "missing_native", + "illegal_window", + "missing_graph", + "wrong_k", + "bad_hash", + "missing_provenance", + ], +) +def test_reuse_keeps_all_structural_guards(tmp_path, fault): + payload = _payload() + if fault == "nonfinite": + payload["cost_tables"]["64"]["step_time_ms"][1] = float("nan") + if fault == "missing_native": + payload["cost_tables"]["64"]["token_counts"][0] = 64 + if fault == "illegal_window": + payload["cost_tables"]["64"]["token_counts"][1] = 257 + if fault == "missing_graph": + payload["cost_tables"] = {"32": payload["cost_tables"]["64"]} + if fault == "wrong_k": + payload["engine_fingerprint"]["max_draft_len"] = 5 + payload["engine_fingerprint_sha256"] = hashlib.sha256( + json.dumps( + payload["engine_fingerprint"], sort_keys=True, separators=(",", ":") + ).encode() + ).hexdigest() + if fault == "bad_hash": + payload["engine_fingerprint_sha256"] = "0" * 64 + if fault == "missing_provenance": + payload["measurements"].pop() + with pytest.raises((TypeError, ValueError)): + _load(tmp_path, payload) + + +def test_provided_drain_policy_prices_and_margin_are_unchanged_in_both_modes(): + kwargs = { + "tables": {64: ExactSpsCostRow((0, 128), (10.0, 8.0))}, + "max_draft_len": 3, + "minimum_predicted_gain": 0.01, + "iteration_drain_guard": ExactSpsDrainGuard(1.0, 1.0, 0.0, 64, "a" * 64), + } + strict = ExactSpsCostTable(**kwargs) + reuse = ExactSpsCostTable(**kwargs, validate_artifact_freshness=False) + assert _select(strict) == _select(reuse) == 0 + assert strict.iteration_drain_guard == reuse.iteration_drain_guard + assert strict.tables == reuse.tables + assert strict.minimum_predicted_gain == reuse.minimum_predicted_gain + assert strict.collective_identity_words != reuse.collective_identity_words + assert len(strict.collective_identity_words) == len(reuse.collective_identity_words) == 8 + + +@pytest.mark.parametrize("strict", [False, True]) +def test_freshness_mode_serialization_and_telemetry_capture_only_the_bool(strict): + class Args(BaseModel): + speculative_config: DSparkDecodingConfig + + config = DSparkDecodingConfig( + max_draft_len=3, + enable_confidence_scheduling=True, + validate_confidence_artifact_freshness=strict, + confidence_sps_table_path="/private/original/sps.json", + confidence_sts_path="/private/original/sts.json", + ) + args = Args.model_validate_json(Args(speculative_config=config).model_dump_json()) + field_path = "speculative_config.validate_confidence_artifact_freshness" + expected = {"path": field_path, "kind": "value", "capture_policy": "bool"} + assert expected in manifest_rows(Args) + captured, metadata = collect_llm_api_config_payloads(args) + assert json.loads(metadata)["capture_succeeded"] is True + assert json.loads(captured)[field_path] is strict + assert "private" not in captured + golden_path = Path(__file__).parents[5] / "tensorrt_llm/usage/llm_args_golden_manifest.json" + assert expected in json.loads(golden_path.read_text(encoding="utf-8"))["TorchLlmArgs"] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph_runtime.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph_runtime.py new file mode 100644 index 000000000000..16cd659ec811 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph_runtime.py @@ -0,0 +1,275 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU tests for DSpark ragged CUDA-graph keys and padding state.""" + +from types import SimpleNamespace + +import numpy as np +import pytest + +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( + CUDA_GRAPH_DUMMY_REQUEST_ID, + ADPShapeAgreement, + CUDAGraphRunner, + KeyType, + cuda_graph_dummy_request_id, +) +from tensorrt_llm._torch.pyexecutor.sampler.sampler_common import SampleType + + +class _Batch: + def __init__(self, generation_requests): + self.generation_requests = list(generation_requests) + self.context_requests = [] + self.encoder_requests = [] + self.can_run_cuda_graph = True + + @property + def batch_size(self): + return len(self.generation_requests) + + @property + def num_context_requests(self): + return 0 + + @property + def num_generation_requests(self): + return len(self.generation_requests) + + +def _runner_stub(): + runner = CUDAGraphRunner.__new__(CUDAGraphRunner) + runner.clear = lambda: None + return runner + + +def _agreement(batch, *, iteration=9, peer_sizes=(3, 4), reuse=True): + return ADPShapeAgreement( + iteration=iteration, + batch_identity=id(batch), + local_batch_size=batch.batch_size, + peer_batch_sizes=peer_sizes, + all_can_graph=True, + widest_batch_size=max(peer_sizes), + graph_batch_size=4, + draft_len=5, + padding_ready=True, + ragged_bucket=16, + reuse_graph_shape=reuse, + ) + + +def test_padding_dummy_ids_are_disjoint(): + ids = { + cuda_graph_dummy_request_id(draft_len, variant=variant, max_draft_len=5) + for variant in (0, 1) + for draft_len in range(6) + } + + assert len(ids) == 12 + assert ( + cuda_graph_dummy_request_id(5, variant=0, max_draft_len=5) + == CUDA_GRAPH_DUMMY_REQUEST_ID - 5 + ) + with pytest.raises(ValueError, match="invalid CUDA padding dummy identity"): + cuda_graph_dummy_request_id(5, variant=2, max_draft_len=5) + + +def test_shape_agreement_is_bound_to_iteration_batch_and_phase(): + batch = _Batch([object()] * 3) + agreement = _agreement(batch) + + assert agreement.matches(batch, 9, padded=False) + assert not agreement.matches(batch, 8, padded=False) + assert not agreement.matches(_Batch([object()] * 3), 9, padded=False) + + batch.generation_requests.append(object()) + assert agreement.matches(batch, 9, padded=True) + assert not agreement.matches(batch, 9, padded=False) + + +def test_ragged_graph_key_requires_every_row_and_the_agreed_bucket(): + runner = _runner_stub() + runner._dspark_trims_submitted_tokens = True + runner.agreed_ragged_bucket = 22 + batch = _Batch( + [ + SimpleNamespace(py_verify_len=3), + SimpleNamespace(py_verify_len=5), + ] + ) + + assert runner._ragged_verify_bucket(batch) == 22 + batch.generation_requests[1].py_verify_len = None + assert runner._ragged_verify_bucket(batch) is None + runner._dspark_trims_submitted_tokens = False + assert runner._ragged_verify_bucket(batch) is None + + +def test_ragged_graph_key_controls_the_submitted_token_count(): + runner = _runner_stub() + + assert ( + runner._get_num_tokens_for_key( + KeyType( + batch_size=4, + draft_len=5, + is_first_draft=False, + ragged_verify_bucket=14, + ) + ) + == 14 + ) + with pytest.raises(ValueError, match="mixed context"): + runner._get_num_tokens_for_key( + KeyType( + batch_size=4, + draft_len=5, + is_first_draft=False, + num_contexts=1, + ragged_verify_bucket=14, + ) + ) + + +@pytest.mark.parametrize( + ("padded_size", "num_requests", "expected"), + [(4, 3, True), (4, 4, True), (8, 3, False), (0, 3, False), (4, 0, False)], +) +def test_will_pad_to_matches_the_size_guards(padded_size, num_requests, expected): + runner = _runner_stub() + runner.supported_batch_sizes = [1, 2, 4] + runner.config = SimpleNamespace(batch_size=4) + + assert runner.will_pad_to(padded_size, num_requests) is expected + + +def test_zero_real_padding_uses_distinct_low_and_high_dummies(): + runner = _runner_stub() + runner.enabled = True + runner.padding_enabled = True + runner.max_supported_batch_size = 4 + runner.config = SimpleNamespace( + enable_attention_dp=False, + mapping=SimpleNamespace(tp_size=1), + batch_size=4, + ) + runner.adp_shape_agreement = None + runner.ragged_pad_verify_len = 2 + runner.ragged_zero_real_high_rows = 2 + runner.spec_config = SimpleNamespace(enable_confidence_scheduling=True, max_draft_len=5) + runner._can_run_cuda_graph_batch = lambda _batch: True + runner._round_up_batch_size_with_draft_len = lambda *_args: 4 + + scheduled_dummy = SimpleNamespace(is_attention_dp_dummy=True, py_verify_len=3) + low = SimpleNamespace(py_verify_len=None) + high = SimpleNamespace(py_verify_len=None) + runner._get_or_create_padding_dummy = ( + lambda _resource, _draft, variant=0: low if variant == 0 else high + ) + batch = _Batch([scheduled_dummy]) + + added = runner._get_padded_batch(batch, SimpleNamespace(), runtime_draft_len=5) + + assert added == 3 + assert batch.generation_requests == [scheduled_dummy, high, low, low] + assert [row.py_verify_len for row in batch.generation_requests] == [3, 3, 2, 2] + assert sum(1 + row.py_verify_len for row in batch.generation_requests) == 14 + + +def test_releasing_padding_dummies_invalidates_the_cached_agreement(): + runner = _runner_stub() + low = object() + high = object() + runner.padding_dummy_requests = {5: low} + runner.secondary_padding_dummy_requests = {5: high} + runner.adp_shape_agreement = object() + freed = [] + runner._padding_dummy_managers = lambda _resource_manager: [ + SimpleNamespace(free_resources=freed.append) + ] + + assert runner.release_padding_dummy(SimpleNamespace(), 5) + assert freed == [low, high] + assert runner.adp_shape_agreement is None + + +def _tier_lookup_runner(batch, *, confidence_enabled, in_graph_sampling, peer_tier): + runner = _runner_stub() + payloads = [] + + def gather(payload): + payloads.append(tuple(payload)) + peer = list(payload) + peer[-1] = peer_tier.value + return np.asarray([payload, peer], dtype=np.int64) + + runner.config = SimpleNamespace( + enable_attention_dp=True, + mapping=SimpleNamespace(tp_size=2), + dist=SimpleNamespace(tp_allgather_int64=gather), + use_mrope=False, + enable_in_graph_sampling=in_graph_sampling, + ) + runner.enabled = True + runner.enable_in_graph_sampling = in_graph_sampling + runner.is_encoder_decoder = False + runner.enable_encoder_decoder_mixed_cuda_graph = False + runner._dspark_confidence_enabled = confidence_enabled + runner._dspark_trims_submitted_tokens = confidence_enabled + runner.adp_shape_agreement = None + runner.agreed_ragged_bucket = 16 if confidence_enabled else None + runner.sparse_config = None + runner._capture_sample_type = None + runner._sample_type_resolver = lambda _batch, _override: SampleType.FAST + runner._capture_allowed = False + runner.graph_metadata = { + KeyType( + batch_size=4, + draft_len=5, + is_first_draft=False, + sample_type=tier, + ragged_verify_bucket=runner.agreed_ragged_bucket, + ): {"attn_metadata": tier, "spec_metadata": None} + for tier in (SampleType.FULL, SampleType.FAST) + } + return runner, payloads + + +@pytest.mark.parametrize("confidence_enabled", [False, True]) +@pytest.mark.parametrize("peer_tier", [SampleType.FAST, SampleType.FULL]) +def test_graph_lookup_agrees_sampling_tier_without_aliasing_ragged_bucket( + confidence_enabled, peer_tier +): + batch = _Batch([SimpleNamespace(py_draft_tokens=[1] * 5, py_verify_len=3) for _ in range(4)]) + runner, payloads = _tier_lookup_runner( + batch, confidence_enabled=confidence_enabled, in_graph_sampling=True, peer_tier=peer_tier + ) + + metadata, _, key = runner.maybe_get_cuda_graph(batch, True, SimpleNamespace()) + + assert metadata is peer_tier + assert key.sample_type is peer_tier + assert payloads == [ + (True, 4, 16, 5, SampleType.FAST.value) + if confidence_enabled + else (True, 4, SampleType.FAST.value) + ] + + +@pytest.mark.parametrize("in_graph_sampling", [False, True]) +def test_cached_shape_does_not_replace_in_graph_sampling_tier_agreement(in_graph_sampling): + batch = _Batch([SimpleNamespace(py_draft_tokens=[1] * 5, py_verify_len=3) for _ in range(4)]) + runner, payloads = _tier_lookup_runner( + batch, + confidence_enabled=True, + in_graph_sampling=in_graph_sampling, + peer_tier=SampleType.FAST, + ) + runner.adp_shape_agreement = _agreement(batch, peer_sizes=(4, 4)) + + metadata, _, key = runner.maybe_get_cuda_graph(batch, True, SimpleNamespace()) + + assert metadata is (SampleType.FAST if in_graph_sampling else SampleType.FULL) + assert key.sample_type is metadata + assert len(payloads) == int(in_graph_sampling) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_select.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_select.py new file mode 100644 index 000000000000..0a1dc79aff7a --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_select.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU tests for fresh-confidence DSpark device-window selection.""" + +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark_device_select import ( + DeviceWindowWorkspace, + gather_packed_draft_tokens, + select_windows_device, +) +from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig + + +def test_workspace_slices_keep_persistent_storage(): + workspace = DeviceWindowWorkspace.allocate(max_rows=4, max_tokens=16, device="cpu") + owners = [ + workspace.verify_lens, + workspace.qo_indptr, + workspace.req_idx, + workspace.kv_correction, + workspace.past_seen, + ] + addresses = [tensor.data_ptr() for tensor in owners] + workspace.validate(device=torch.device("cpu"), num_rows=2, num_tokens=8) + workspace.validate(device=torch.device("cpu"), num_rows=4, num_tokens=16) + assert [tensor.data_ptr() for tensor in owners] == addresses + + +@pytest.mark.parametrize("num_rows,num_tokens", [(5, 16), (4, 17)]) +def test_workspace_rejects_out_of_capacity_shape(num_rows, num_tokens): + workspace = DeviceWindowWorkspace.allocate(max_rows=4, max_tokens=16, device="cpu") + with pytest.raises(ValueError, match="capacity"): + workspace.validate(device=torch.device("cpu"), num_rows=num_rows, num_tokens=num_tokens) + + +def test_workspace_rejects_wrong_row_extent_dtype(): + workspace = DeviceWindowWorkspace.allocate(max_rows=4, max_tokens=16, device="cpu") + workspace.verify_lens = workspace.verify_lens.to(torch.int64) + with pytest.raises(TypeError, match="verify_lens must have dtype"): + workspace.validate(device=torch.device("cpu"), num_rows=4, num_tokens=16) + + +def test_fresh_confidence_spends_budget_on_the_best_real_row(): + result = select_windows_device( + confidence_logits=torch.tensor( + [[10.0, 10.0, 10.0, 10.0, 10.0], [-10.0, -10.0, -10.0, -10.0, -10.0]] + ), + slot_idx=torch.tensor([0, 1, 0, 0]), + num_real=torch.tensor(2), + budget=torch.tensor(2), + graph_num_tokens=8, + cfg=DSparkScheduleConfig(block_size=5, min_verify_len=1), + pad_len=1, + ) + + assert result.verify_lens.tolist() == [4, 2, 1, 1] + assert result.qo_indptr.tolist() == [0, 4, 6, 7, 8] + assert result.req_idx.tolist() == [0, 0, 0, 0, 1, 1, 2, 3] + assert result.kv_correction.tolist() == [-3, -2, -1, 0, -1, 0, 0, 0] + + +def test_stale_row_fails_open_to_neutral_confidence(): + result = select_windows_device( + confidence_logits=torch.full((2, 5), -10.0), + slot_idx=torch.tensor([0, 1]), + num_real=torch.tensor(2), + budget=torch.tensor(1), + graph_num_tokens=5, + cfg=DSparkScheduleConfig(block_size=5, min_verify_len=1), + stamp=torch.tensor([7, 6]), + expected_stamp=torch.tensor(7), + ) + + assert result.verify_lens.tolist() == [2, 3] + + +def test_gather_packed_draft_tokens_excludes_each_bonus(): + packed = gather_packed_draft_tokens( + next_draft_tokens=torch.tensor([[10, 11, 12], [20, 21, 22]]), + batch_slots=torch.tensor([1, 0]), + verify_lens=torch.tensor([3, 2], dtype=torch.int32), + qo_indptr=torch.tensor([0, 3, 5], dtype=torch.int32), + num_real=2, + total_draft_tokens=3, + ) + + assert packed.tolist() == [20, 21, 10] + + +def test_gather_packed_draft_tokens_rejects_negative_size(): + with pytest.raises(ValueError, match="non-negative"): + gather_packed_draft_tokens( + next_draft_tokens=torch.empty((1, 1)), + batch_slots=torch.tensor([0]), + verify_lens=torch.tensor([1], dtype=torch.int32), + qo_indptr=torch.tensor([0, 1], dtype=torch.int32), + num_real=1, + total_draft_tokens=-1, + ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_dummy_slots_cpu.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_dummy_slots_cpu.py new file mode 100644 index 000000000000..7b5fecb3d949 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_dummy_slots_cpu.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Host slot-ownership regression with CPU tensors in place of CUDA allocations.""" + +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm._torch.speculative.dspark as dspark +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import cuda_graph_dummy_request_id +from tensorrt_llm._torch.speculative.dspark import DSparkSpecMetadata, DSv4DSparkWorker + +pytestmark = pytest.mark.cpu_only + + +@pytest.mark.parametrize("physical_k", range(1, 9)) +def test_all_padding_families_preserve_real_slot_ownership(monkeypatch, physical_k): + # Exercise the real initialization and host prepare methods. This deliberately + # does not validate device copies, graph replay, or target KV-cache contents. + original_zeros = torch.zeros + original_empty = torch.empty + original_full = torch.full + + def cpu_zeros(*args, **kwargs): + kwargs["device"] = "cpu" + return original_zeros(*args, **kwargs) + + def cpu_empty(*args, **kwargs): + kwargs["device"] = "cpu" + return original_empty(*args, **kwargs) + + def cpu_full(*args, **kwargs): + kwargs["device"] = "cpu" + return original_full(*args, **kwargs) + + monkeypatch.setattr(torch, "zeros", cpu_zeros) + monkeypatch.setattr(torch, "empty", cpu_empty) + monkeypatch.setattr(torch, "full", cpu_full) + monkeypatch.setattr(dspark, "prefer_pinned", lambda: False) + worker = object.__new__(DSv4DSparkWorker) + worker.spec_config = SimpleNamespace(max_draft_len=physical_k) + worker._win_inited = False + worker._attention_warmup_attempted = True + worker.return_confidence = False + worker._confidence_logits = None + worker._mixed_confidence_tracking = False + worker._confidence_incarnations = {} + draft = SimpleNamespace( + block_size=physical_k, + num_stages=1, + _attn_params={"window_size": 4, "head_dim": 2}, + ) + DSv4DSparkWorker._lazy_init(worker, draft, SimpleNamespace(max_num_requests=4)) + meta = SimpleNamespace( + _dspark_worker=worker, batch_indices_cuda=torch.empty(4, dtype=torch.int32) + ) + for variant in (0, 1): + for runtime_k in range(physical_k + 1): + dummy = cuda_graph_dummy_request_id( + runtime_k, variant=variant, max_draft_len=physical_k + ) + # Include ADP idle, the secondary zero-real family, and a live row. + meta.request_ids = [1000, 0, dummy] + meta.num_generations = 3 + DSparkSpecMetadata.prepare(meta) + assert worker._req_to_slot == {1000: 0} + assert list(worker._free_slots) == [1, 2, 3] + assert worker._batch_to_slot[:3].tolist() == [0, 4, 4] + # Removing a real request releases exactly its slot, never a dummy's. + meta.request_ids = [dummy, 0] + meta.num_generations = 2 + DSparkSpecMetadata.prepare(meta) + assert worker._req_to_slot == {} + assert sorted(worker._free_slots) == [0, 1, 2, 3] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_embedded_confidence_loading.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_embedded_confidence_loading.py new file mode 100644 index 000000000000..01881356f6e9 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_embedded_confidence_loading.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Embedded-wrapper validation through the real shared parameter walk.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch.models import modeling_dspark, modeling_dspark_v41 +from tensorrt_llm._torch.models.modeling_speculative import DSparkConfidenceHead + +pytestmark = pytest.mark.cpu_only + + +def _wrapper(monkeypatch, v41): + module = modeling_dspark_v41 if v41 else modeling_dspark + cls = module.DSv41DSparkForCausalLM if v41 else module.DSv4DSparkForCausalLM + wrapper = object.__new__(cls) + nn.Module.__init__(wrapper) + stage = nn.Module() + stage.main_proj = nn.Linear(2, 2, bias=False) + stage.main_norm = nn.RMSNorm(2) + stage.norm = nn.RMSNorm(2) + stage.markov_head = None + stage.confidence_head = DSparkConfidenceHead(hidden_size=2, block_size=2) + model = nn.Module() + model.mtp_layers = nn.ModuleList([stage]) + model.config = SimpleNamespace( + q_lora_rank=1, + num_attention_heads=1, + num_hidden_layers=0, + num_nextn_predict_layers=1, + kv_lora_rank=1, + ) + model.model_config = SimpleNamespace( + mapping=SimpleNamespace( + tp_rank=0, + tp_size=1, + cp_rank=0, + cp_size=1, + enable_attention_dp=False, + ) + ) + model.post_load_weights = Mock() + model.cache_attn_weights_from_state_dict = Mock() + wrapper.dspark_model = model + wrapper.config = model.config + wrapper.num_stages = 1 + # Remapping has separate checkpoint-layout tests. This test starts with + # model-named tensors and does not mock the loader or its parameter copies. + remapper = "remap_dspark_v41_draft_keys" if v41 else "remap_dspark_draft_keys" + monkeypatch.setattr(module, remapper, lambda weights, *args, **kwargs: weights) + weights = {name: torch.full_like(param, 2.0) for name, param in model.named_parameters()} + return wrapper, weights + + +@pytest.mark.parametrize("v41", [False, True]) +@pytest.mark.parametrize("bias", [False, True]) +def test_embedded_confidence_load_validates_and_copies_real_parameters(monkeypatch, v41, bias): + wrapper, weights = _wrapper(monkeypatch, v41) + head = wrapper.dspark_model.mtp_layers[-1].confidence_head + prefix = "mtp_layers.0.confidence_head." + if bias: + weights[prefix + "proj.bias"] = torch.tensor([3.0]) + weight_pointer = head.proj.weight.data_ptr() + rng_before = torch.get_rng_state() + wrapper.load_weights(weights) + assert head.proj.weight.data_ptr() == weight_pointer + assert torch.equal(torch.get_rng_state(), rng_before) + assert torch.equal(head.proj.weight, weights[prefix + "proj.weight"]) + assert (head.proj.bias is not None) is bias + if bias: + assert torch.equal(head.proj.bias, weights[prefix + "proj.bias"]) + wrapper.dspark_model.post_load_weights.assert_called_once() + + +@pytest.mark.parametrize("v41", [False, True]) +@pytest.mark.parametrize("fault", ["unknown", "missing_weight", "broadcast_weight"]) +def test_embedded_confidence_rejects_incomplete_or_unknown_tensors(monkeypatch, v41, fault): + wrapper, weights = _wrapper(monkeypatch, v41) + prefix = "mtp_layers.0.confidence_head." + if fault == "unknown": + weights[prefix + "bias"] = torch.tensor([99.0]) + elif fault == "missing_weight": + weights.pop(prefix + "proj.weight") + weights[prefix + "proj.bias"] = torch.tensor([3.0]) + else: + weights[prefix + "proj.weight"] = torch.tensor([[2.0]]) + with pytest.raises(ValueError, match="confidence head checkpoint"): + wrapper.load_weights(weights) + wrapper.dspark_model.post_load_weights.assert_not_called() + + +def test_v41_entirely_absent_confidence_head_remains_supported(monkeypatch): + wrapper, weights = _wrapper(monkeypatch, True) + weights.pop("mtp_layers.0.confidence_head.proj.weight") + wrapper.load_weights(weights) + assert wrapper.dspark_model.mtp_layers[-1].confidence_head is None + wrapper.dspark_model.post_load_weights.assert_called_once() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_executor_policy.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_executor_policy.py new file mode 100644 index 000000000000..f7d80b2e0110 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_executor_policy.py @@ -0,0 +1,326 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only tests for DSpark policy agreement in the executor.""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ADPShapeAgreement +from tensorrt_llm._torch.pyexecutor.py_executor import ( + _DSPARK_ADP_WIRE_TRAILER_LEN, + _DSPARK_EXACT_NATIVE_YIELD_INDEX, + _DSPARK_EXACT_WIRE_PREFIX_LEN, + PyExecutor, + _classify_dspark_exact_generation_rows, + _decode_dspark_exact_expected_yields, + _dspark_exact_common_graph_batch_size, + _dspark_exact_secondary_padding_ready, + _encode_dspark_exact_expected_yields, + _validate_dspark_adp_acceptance_gate, + _validate_dspark_exact_bucket, +) + + +class _Batch: + def __init__(self, generation_requests): + self.generation_requests = generation_requests + self.context_requests = [] + + @property + def batch_size(self): + return len(self.generation_requests) + + @property + def num_context_requests(self): + return 0 + + @property + def num_generation_requests(self): + return len(self.generation_requests) + + +def _request(*, adp_dummy=False, other_dummy=False): + return SimpleNamespace( + is_dummy=bool(adp_dummy or other_dummy), + is_attention_dp_dummy=bool(adp_dummy), + ) + + +def _payload(*, rows, native, compact): + payload = [0] * (_DSPARK_EXACT_WIRE_PREFIX_LEN + len(compact) + _DSPARK_ADP_WIRE_TRAILER_LEN) + payload[1] = rows + payload[_DSPARK_EXACT_NATIVE_YIELD_INDEX] = native + payload[_DSPARK_EXACT_WIRE_PREFIX_LEN : _DSPARK_EXACT_WIRE_PREFIX_LEN + len(compact)] = compact + return payload + + +def _agreement(batch, peer_sizes=(3, 4)): + return ADPShapeAgreement( + iteration=9, + batch_identity=id(batch), + local_batch_size=batch.batch_size, + peer_batch_sizes=peer_sizes, + all_can_graph=True, + widest_batch_size=max(peer_sizes), + graph_batch_size=4, + draft_len=5, + padding_ready=True, + ragged_bucket=16, + reuse_graph_shape=True, + ) + + +def test_exact_yield_encoding_is_conservative_and_pads_missing_cells(): + exact_local = SimpleNamespace( + native_expected_yield=10.0, + compact_expected_yields=(9.25, float("nan")), + ) + + native, compact = _encode_dspark_exact_expected_yields( + exact_local=exact_local, num_cells=3, yield_scale=100 + ) + + assert native == 1000 + assert compact == (925, -1, -1) + + +def test_exact_yield_encoding_rejects_an_invalid_scale(): + with pytest.raises(ValueError, match="scale must be positive"): + _encode_dspark_exact_expected_yields(exact_local=None, num_cells=1, yield_scale=0) + + +def test_exact_yield_decode_aggregates_yield_and_worst_rank_loss(): + cells = ((128, 512), (128, 640)) + payloads = [ + _payload( + rows=2, + native=10_000_000, + compact=(8_000_000, 9_000_000), + ), + _payload( + rows=4, + native=20_000_000, + compact=(18_000_000, 14_000_000), + ), + _payload(rows=0, native=0, compact=(0, 0)), + ] + + native, compact, max_loss = _decode_dspark_exact_expected_yields( + payloads=payloads, + exact_cells=cells, + graph_batch_size=128, + yield_scale=1_000_000, + ) + + assert native == 30.0 + assert compact == {512: 26.0, 640: 23.0} + assert max_loss == {512: 1.0, 640: 1.5} + + +def test_exact_yield_decode_fails_closed_on_an_infeasible_peer(): + payloads = [ + _payload(rows=2, native=10_000_000, compact=(8_000_000,)), + _payload(rows=4, native=20_000_000, compact=(-1,)), + ] + + _, compact, max_loss = _decode_dspark_exact_expected_yields( + payloads=payloads, + exact_cells=((128, 512),), + graph_batch_size=128, + yield_scale=1_000_000, + ) + + assert compact == {512: 0.0} + assert max_loss == {512: 0.0} + + +def test_exact_row_classifier_authenticates_only_the_idle_adp_dummy(): + real = [_request(), _request()] + assert _classify_dspark_exact_generation_rows(real) == (real, False) + assert _classify_dspark_exact_generation_rows([]) == (None, False) + assert _classify_dspark_exact_generation_rows([_request(adp_dummy=True)]) == ( + [], + True, + ) + assert _classify_dspark_exact_generation_rows([real[0], _request(adp_dummy=True)]) == ( + None, + False, + ) + assert _classify_dspark_exact_generation_rows([_request(other_dummy=True)]) == (None, False) + + +def test_exact_common_graph_size_preserves_all_rank_termination(): + def round_up(rows): + return 16 if rows <= 16 else 32 + + assert _dspark_exact_common_graph_batch_size([[0, 16], [0, 0]], round_up) == 16 + + def reject_round_up(_rows): + raise AssertionError("all-zero peers must not perform a graph lookup") + + assert _dspark_exact_common_graph_batch_size([[0, 0], [0, 0]], reject_round_up) == 0 + + +@pytest.mark.parametrize( + ("budget", "secondary_ready", "expected"), + [(48, False, True), (49, False, True), (50, False, False), (50, True, True)], +) +def test_secondary_dummy_is_required_only_for_a_multirow_remainder( + budget, secondary_ready, expected +): + payloads = [[0, 16, 0], [0, 0, int(secondary_ready)]] + assert ( + _dspark_exact_secondary_padding_ready( + payloads, + graph_batch_size=16, + verifier_budget=budget, + secondary_ready_index=2, + ) + is expected + ) + + +def test_exact_bucket_must_match_the_selected_graph_key(): + _validate_dspark_exact_bucket(exact_shape=(16, 48, 3), bucket=48) + _validate_dspark_exact_bucket(exact_shape=None, bucket=None) + with pytest.raises(RuntimeError, match="selected captured verifier bucket"): + _validate_dspark_exact_bucket(exact_shape=(16, 48, 3), bucket=47) + + +def test_attention_dp_rejects_a_rank_local_acceptance_gate(): + with pytest.raises(ValueError, match="acceptance_rate_window_size"): + _validate_dspark_adp_acceptance_gate( + confidence_enabled=True, + attention_dp_enabled=True, + speculation_gate_enabled=True, + ) + _validate_dspark_adp_acceptance_gate( + confidence_enabled=True, + attention_dp_enabled=False, + speculation_gate_enabled=True, + ) + + +@pytest.mark.parametrize( + ("peer_sizes", "expected"), + [((3, 4), (True, True)), ((3, 0), (False, True))], +) +def test_can_queue_reuses_policy_agreement_without_another_collective(peer_sizes, expected): + batch = _Batch([object()] * 3) + + def reject_allgather(_value): + raise AssertionError("cached agreement issued another collective") + + executor = PyExecutor.__new__(PyExecutor) + executor.enable_attention_dp = True + executor.dist = SimpleNamespace(tp_size=8, tp_allgather=reject_allgather) + executor.model_engine = SimpleNamespace( + _dspark_confidence_enabled=True, + cuda_graph_runner=SimpleNamespace(adp_shape_agreement=_agreement(batch, peer_sizes)), + ) + executor.kv_connector_manager = None + executor.kv_cache_transceiver = None + executor.drafter = None + executor.iter_counter = 9 + executor._dspark_dynamic_handled_signature = (9, id(batch), 3, 0, 3) + executor._handle_dynamic_draft_len = lambda _batch: None + + assert executor._can_queue(batch) == expected + + +def _draft_len_executor(*, confidence_enabled): + spec_mode = SimpleNamespace( + support_dynamic_draft_len=lambda: False, + is_pard=lambda: False, + ) + spec_config = SimpleNamespace( + draft_len_schedule=None, + spec_dec_mode=spec_mode, + enable_confidence_scheduling=confidence_enabled, + use_rejection_sampling=False, + is_linear_tree=True, + ) + executor = PyExecutor.__new__(PyExecutor) + executor.model_engine = SimpleNamespace( + max_draft_len=5, + max_total_draft_tokens=5, + runtime_draft_len=None, + spec_config=spec_config, + cuda_graph_runner=SimpleNamespace(adp_shape_agreement=None), + ) + executor.speculation_permanently_disabled = False + executor.iter_counter = 9 + return executor + + +def test_fixed_draft_len_keeps_static_path_without_dynamic_signature(): + executor = _draft_len_executor(confidence_enabled=False) + batch = _Batch([object()]) + + executor._handle_dynamic_draft_len(batch) + executor._handle_dynamic_draft_len(batch) + + assert executor.model_engine.runtime_draft_len == 5 + assert not hasattr(executor, "_dspark_dynamic_handled_signature") + + +def test_confidence_draft_len_runs_only_once_per_iteration(): + executor = _draft_len_executor(confidence_enabled=True) + request = SimpleNamespace( + py_draft_tokens=[1, 2, 3, 4, 5], + py_needs_onehot_draft_probs=False, + ) + batch = _Batch([request]) + calls = 0 + + def select_draft_len(_batch): + nonlocal calls + calls += 1 + return 3 + + executor._dspark_confidence_draft_len = select_draft_len + executor._handle_dynamic_draft_len(batch) + executor._handle_dynamic_draft_len(batch) + + assert calls == 1 + assert executor.model_engine.runtime_draft_len == 3 + assert request.py_draft_tokens == [1, 2, 3] + assert executor._dspark_dynamic_handled_signature == (9, id(batch), 1, 0, 1) + + +def test_rebalance_suspends_both_dummy_variants_and_releases_the_pair_once(): + low = SimpleNamespace(py_request_id=10) + high = SimpleNamespace(py_request_id=11) + suspended_ids = [] + resumed_ids = [] + + class _Manager: + def is_request_active(self, request_id): + return request_id in {10, 11} + + def suspend_request(self, request): + suspended_ids.append(request.py_request_id) + + def resume_request(self, request): + resumed_ids.append(request.py_request_id) + return request.py_request_id != 11 + + released = [] + runner = SimpleNamespace( + padding_dummy_requests={5: low}, + secondary_padding_dummy_requests={5: high}, + release_padding_dummy=lambda resource, draft_len: released.append((resource, draft_len)), + ) + executor = PyExecutor.__new__(PyExecutor) + executor.model_engine = SimpleNamespace(cuda_graph_runner=runner) + executor.resource_manager = object() + manager = _Manager() + + suspended = executor._suspend_padding_dummies_for_rebalance(manager) + executor._resume_padding_dummies_after_rebalance(manager, suspended) + + assert suspended == [(5, 0, low), (5, 1, high)] + assert suspended_ids == [10, 11] + assert resumed_ids == [10, 11] + assert released == [(executor.resource_manager, 5)] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_fused_attention_route.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_fused_attention_route.py new file mode 100644 index 000000000000..7b1a6e030a4b --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_fused_attention_route.py @@ -0,0 +1,328 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU-only DSpark routing contracts; no Torch, CuTe, or CUDA runtime is loaded. + +Execute the production Python guards with tensor metadata and mocked launchers. +The AST extraction follows the existing K2/K3 routing-test prototype. Scalar +mask checks are not a substitute for compiling and testing the GPU kernels. +""" + +import ast +import math +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +_ROOT = Path(__file__).resolve().parents[5] +_OPS = _ROOT / "tensorrt_llm/_torch/custom_ops" +_ATTENTION = _OPS / "dspark_attention_custom_op.py" +_PREPARATION = _OPS / "dspark_rmsnorm_rope_custom_op.py" +_KERNELS = _ROOT / "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark" +_MODEL = _ROOT / "tensorrt_llm/_torch/models/modeling_dspark.py" +_PHYSICAL_KS = tuple(range(1, 9)) + + +def _tree(path): + return ast.parse(path.read_text(encoding="utf-8")) + + +def _load(path, name, namespace, owner=None): + nodes = _tree(path).body + if owner is not None: + nodes = next(node for node in nodes if getattr(node, "name", None) == owner).body + node = next(node for node in nodes if getattr(node, "name", None) == name) + node.decorator_list = [] + future = ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0) + module = ast.fix_missing_locations(ast.Module(body=[future, node], type_ignores=[])) + exec(compile(module, str(path), "exec"), namespace) + return namespace[name] + + +def _constants(path): + return { + node.targets[0].id: ast.literal_eval(node.value) + for node in _tree(path).body + if isinstance(node, ast.Assign) + and isinstance(node.targets[0], ast.Name) + and node.targets[0].id.startswith("_DSV4_DSPARK_") + } + + +def _tensor(shape, dtype="bf16", *, cuda=True, contiguous=True, strides=None): + if strides is None: + strides = tuple(math.prod(shape[i + 1 :]) for i in range(len(shape))) + return SimpleNamespace( + shape=shape, + ndim=len(shape), + dtype=dtype, + is_cuda=cuda, + is_contiguous=lambda: contiguous, + numel=lambda: math.prod(shape), + stride=lambda dim=None: strides if dim is None else strides[dim], + data_ptr=lambda: 0, + ) + + +def _guards(sm=100): + namespace = { + "torch": SimpleNamespace(bfloat16="bf16", float32="fp32", int32="i32", int64="i64"), + "get_sm_version": lambda: sm, + "_log_unsupported": lambda *args: False, + } + namespace.update(_constants(_PREPARATION)) + for name in ( + "_get_dspark_arch_str", + "is_fused_dspark_rmsnorm_rope_supported", + "_is_dspark_cache_write_supported", + "_is_dspark_draft_block_supported", + "is_fused_dspark_attention_preparation_supported", + ): + _load(_PREPARATION, name, namespace) + namespace.update(_constants(_ATTENTION)) + for name in ( + "is_dsv4_dspark_attention_config_supported", + "is_fused_dsv4_dspark_attention_supported", + "_is_fused_dsv4_dspark_attention_input_supported", + ): + _load(_ATTENTION, name, namespace) + return namespace + + +def _inputs(k, batch=2): + return [ + _tensor((batch, k, 128, 512)), + _tensor((batch, 8, 512)), + _tensor((batch + 1, 128, 512), contiguous=False, strides=(196608, 512, 1)), + _tensor((batch,), "i32"), + _tensor((batch,), "i32"), + _tensor((batch,), "i64"), + _tensor((128,), "fp32"), + _tensor((batch, k, 32, 2), "fp32"), + ] + + +class TestDSparkFusedAttentionRoute(unittest.TestCase): + def test_storage_bounded_range_is_consistent(self): + for path in (_ATTENTION, _PREPARATION): + constants = _constants(path) + self.assertEqual(constants["_DSV4_DSPARK_BLOCK_SIZES"], _PHYSICAL_KS) + self.assertEqual(constants["_DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE"], 8) + kernel = next( + node + for node in _tree(_KERNELS / "attention.py").body + if isinstance(node, ast.ClassDef) and node.name == "DSparkAttention" + ) + page_size = next( + ast.literal_eval(node.value) + for node in kernel.body + if isinstance(node, ast.Assign) and node.targets[0].id == "page_size_draft" + ) + self.assertEqual(_PHYSICAL_KS, tuple(range(1, page_size + 1))) + + def test_config_routes_every_k_and_rejects_out_of_range(self): + supported = _guards()["is_dsv4_dspark_attention_config_supported"] + for k in (-1, 0, *_PHYSICAL_KS, 9, 16, 128): + with self.subTest(k=k): + self.assertEqual(supported(k, 128, 512, 128), k in _PHYSICAL_KS) + for k in _PHYSICAL_KS: + self.assertFalse(supported(k, 64, 512, 128)) + self.assertFalse(supported(k, 128, 256, 128)) + self.assertFalse(supported(k, 128, 512, 64)) + + def test_kernel_constructor_keeps_geometry_and_storage_guards(self): + class KernelBase: + def __init__(self, *args, **kwargs): + self.seq_len_q = kwargs["seq_len_q"] + + namespace = { + "cutlass": SimpleNamespace(Float32="fp32"), + "cute": SimpleNamespace(jit=lambda f: f), + "DSparkAttentionKernel": KernelBase, + } + kernel = _load(_KERNELS / "attention.py", "DSparkAttention", namespace) + args = ("fp32", (128, 128), (128, 256), 1, 8, 128, 0.0) + for k in _PHYSICAL_KS: + with self.subTest(k=k): + op = kernel(*args, arch_str="sm_100", seq_len_q=k, inverse_rope_dim=64) + self.assertEqual(op.seq_len_q, k) + self.assertEqual(op.fixed_cache_seq_len, 128 + k) + self.assertEqual(op.tma_page_size_draft, 128) + for k in (-1, 0, 9, 16): + with self.subTest(k=k), self.assertRaises(ValueError): + kernel(*args, arch_str="sm_100", seq_len_q=k) + with self.assertRaises(ValueError): + kernel(*args, arch_str="sm_100", seq_len_q=8, inverse_rope_dim=32) + with self.assertRaises(ValueError): + kernel(*args[:4], 16, *args[5:], arch_str="sm_100", seq_len_q=8) + + def test_attention_tensor_contract_for_every_k_and_architecture(self): + for sm in (90, 100, 101, 103, 109, 120, 121): + gate = _guards(sm)["_is_fused_dsv4_dspark_attention_input_supported"] + for k in (0, *_PHYSICAL_KS, 9, 16): + with self.subTest(sm=sm, k=k): + self.assertEqual(gate(*_inputs(k)), sm in (100, 103) and k in _PHYSICAL_KS) + + def test_attention_dtype_shape_and_layout_guards_survive(self): + gate = _guards()["_is_fused_dsv4_dspark_attention_input_supported"] + bad_inputs = ( + (0, _tensor((2, 4, 128, 512), dtype="fp32")), + (0, _tensor((2, 4, 64, 512))), + (0, _tensor((2, 4, 128, 256))), + (0, _tensor((2, 4, 512))), + (1, _tensor((2, 7, 512))), + (1, _tensor((2, 9, 512))), + (1, _tensor((2, 8, 512), dtype="fp32")), + (2, _tensor((3, 64, 512))), + (2, _tensor((3, 128, 512), dtype="fp32")), + (2, _tensor((3, 128, 512), strides=(0, 512, 1))), + (2, _tensor((3, 128, 512), strides=(196609, 512, 1))), + (2, _tensor((3, 128, 512), strides=(196608, 513, 1))), + (2, _tensor((3, 128, 512), strides=(196608, 512, 2))), + (3, _tensor((2,), "i64")), + (4, _tensor((2,), "i64")), + (4, _tensor((3,), "i32")), + (5, _tensor((2,), "i32")), + (5, _tensor((3,), "i64")), + (6, _tensor((128,), "bf16")), + (6, _tensor((64,), "fp32")), + (7, _tensor((2, 4, 31, 2), "fp32")), + ) + for index, bad in bad_inputs: + with self.subTest(index=index, shape=bad.shape, dtype=bad.dtype): + inputs = _inputs(4) + inputs[index] = bad + self.assertFalse(gate(*inputs)) + for k in _PHYSICAL_KS: + for index in range(8): + inputs = _inputs(k) + inputs[index].is_cuda = False + self.assertFalse(gate(*inputs)) + if index != 2: # The worker's strided cache-page view is supported. + inputs = _inputs(k) + inputs[index].is_contiguous = lambda: False + self.assertFalse(gate(*inputs)) + + def test_preparation_bounds_and_guards_for_every_k(self): + gate = _guards()["is_fused_dspark_attention_preparation_supported"] + for k in (0, *_PHYSICAL_KS, 9, 16): + with self.subTest(k=k): + inputs = [ + _tensor((2, 1, 512)), + _tensor((2, k, 512)), + _tensor((512,)), + _tensor((2, 32, 2), "fp32"), + _tensor((2 * k, 32, 2), "fp32"), + _inputs(k)[2], + _tensor((2,), "i64"), + _tensor((2,), "i64"), + ] + self.assertEqual(gate(*inputs), k in _PHYSICAL_KS) + if k in _PHYSICAL_KS: + for index in range(len(inputs)): + original = inputs[index].is_cuda + inputs[index].is_cuda = False + self.assertFalse(gate(*inputs)) + inputs[index].is_cuda = original + inputs[4] = _tensor((2 * k + 1, 32, 2), "fp32") + self.assertFalse(gate(*inputs)) + + def test_masks_cover_short_full_wrapped_and_empty_windows(self): + namespace = { + "cutlass": SimpleNamespace(Int32=int, const_expr=lambda x: x), + "cute": SimpleNamespace(elem_less=lambda a, b: a < b), + } + path = _KERNELS / "attention_kernel.py" + mask = _load(path, "is_score_valid", namespace, "DSparkAttentionKernel") + valid_length = _load(path, "get_window_valid_len", namespace, "DSparkAttentionKernel") + kernel = SimpleNamespace( + mma_qk_tiler=(128, 128, 128), + window_valid_len_from_tensor=True, + tma_page_size_win=128, + ) + for k in _PHYSICAL_KS: + for end in (0, 5, 127, 128, 257, 390): + for valid in (-1, 0, 1, 3, 128, 129): + with self.subTest(k=k, end=end, valid=valid): + length = valid_length(kernel, [valid], 0) + expected = {(end - age) % 128 for age in range(min(max(valid, 0), 128))} + actual = { + col for col in range(128) if mask(kernel, col, 0, 128 + k, length, end) + } + self.assertEqual(actual, expected) + draft = { + col for col in range(128) if mask(kernel, col, 1, 128 + k, length, end) + } + self.assertEqual(draft, set(range(k))) + + def test_cache_keys_separate_k_and_arch_but_not_runtime_shapes(self): + namespace = _guards() + cache = {} + compile_kernel = Mock(side_effect=lambda *args: Mock()) + namespace["_dspark_attention_kernel_cache"] = cache + namespace["_compile_dspark_attention"] = compile_kernel + capturing = Mock(return_value=False) + namespace["torch"].cuda = SimpleNamespace(is_current_stream_capturing=capturing) + namespace["torch"].empty_like = lambda q: _tensor(q.shape) + run = _load(_ATTENTION, "_run_dspark_attention", namespace) + for arch in ("sm_100", "sm_103"): + namespace["_get_dspark_arch_str"] = lambda arch=arch: arch + for k in _PHYSICAL_KS: + with self.subTest(k=k, arch=arch): + for batch, scale in ((1, 0.125), (3, 0.25)): + self.assertEqual(run(*_inputs(k, batch), scale).shape, (batch, k, 128, 512)) + self.assertEqual(cache[(k, arch)].call_count, 2) + self.assertEqual(compile_kernel.call_count, 16) + self.assertEqual(capturing.call_count, 16) + self.assertEqual(len({id(compiled) for compiled in cache.values()}), 16) + + def test_graph_capture_only_uses_warmed_k(self): + for k in _PHYSICAL_KS: + namespace = _guards() + namespace["_dspark_attention_kernel_cache"] = {} + namespace["_compile_dspark_attention"] = Mock(side_effect=AssertionError("compile")) + namespace["torch"].cuda = SimpleNamespace(is_current_stream_capturing=lambda: True) + namespace["torch"].empty_like = Mock(side_effect=AssertionError("allocation")) + run = _load(_ATTENTION, "_run_dspark_attention", namespace) + with self.subTest(k=k), self.assertRaisesRegex(RuntimeError, "warmed up"): + run(*_inputs(k), 0.125) + + def test_unsupported_k_warmup_does_not_allocate(self): + namespace = _guards() + warmup = _load(_ATTENTION, "warmup_fused_dsv4_dspark_attention", namespace) + # The metadata-only torch object cannot allocate; early return is required. + for k in (-1, 0, 9, 16): + warmup(k, 1e-6) + + def test_query_storage_and_fallback_routes_remain_separate(self): + attention = ast.unparse(_tree(_KERNELS / "attention.py")) + self.assertIn("(self.num_heads, self.head_dim, self.seq_len_q, batch_size)", attention) + self.assertIn("batch_size * self.seq_len_q * self.inverse_rope_dim", attention) + self.assertIn("self.fixed_cache_seq_len = self.window_size + seq_len_q", attention) + model = {} + _load(_MODEL, "dspark_attention_forward_batched", model) + function = next( + node + for node in _tree(_MODEL).body + if getattr(node, "name", None) == "dspark_attention_forward_batched" + ) + route = next( + node + for node in function.body + if isinstance(node, ast.If) + and ast.unparse(node.test) == "use_fused_dsv4_dspark_attention" + ) + fused = ast.unparse(ast.Module(body=route.body, type_ignores=[])) + fallback = ast.unparse(ast.Module(body=route.orelse, type_ignores=[])) + self.assertIn("cute_dsl_dspark_rmsnorm_rope_draft_block(", fused) + self.assertIn("fused_dsv4_dspark_attention(", fused) + self.assertNotIn("fused_dsv4_dspark_attention(", fallback) + self.assertIn( + "get_dspark_topk_idxs_batched(window_size, block, start_pos, valid_len)", fallback + ) + self.assertIn("torch.cat([cache_rows, kv], dim=1)", fallback) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_evidence.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_evidence.py new file mode 100644 index 000000000000..3fca02f00a02 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_evidence.py @@ -0,0 +1,458 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only, deliberately synthetic protocol fixtures; no production evidence.""" + +import copy +import dataclasses +import hashlib +import json +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace + +from tensorrt_llm._torch.speculative import dspark_mixed_evidence as mixed +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig, DSparkMixedEvidenceConfig + + +def sha(raw): + return hashlib.sha256(raw).hexdigest() + + +def encode(value): + return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() + + +def identity(): + return mixed.MixedRuntimeIdentity(*("abcdef"[index] * 64 for index in range(6))) + + +def geometry(): + # Rank 0 prefills 64 rows; rank 1 decodes with native 32 versus compact 16. + # Every bootstrap/dummy generation is still charged its full K+1 rows. + features = ( + (64, 4096, 8192, 0, 0, 0, 0, 1, 0, 1, 0, 0), + (0, 0, 0, 0, 1024, 128, 0, 0, 8, 0, 8, 0), + ) + ((0,) * 12,) * 6 + return mixed.MixedGeometry(features, 96, (64, 16, 0, 0, 0, 0, 0, 0), (64, 32, 0, 0, 0, 0, 0, 0)) + + +def fixture(): + live = identity() + table = { + "schema": "dspark-mixed-breakable-costs-v1", + "identity": live.payload(), + "protocol_sha256": "1" * 64, + "action_semantics": mixed._ACTION, + "pricing_semantics": mixed._PRICE_SEMANTICS, + "synthetic": False, + "provenance": "closed_measured_breakable_mixed", + "minimum_predicted_gain": 0.01, + "cells": [ + { + "geometry": geometry().payload(), + "native_lower_ms": 7.0, + "candidate_upper_ms": 6.9, + "measurement_result_sha256": "2" * 64, + } + ], + } + receipt = { + "schema": "dspark-mixed-breakable-admission-v1", + "identity": live.payload(), + "costs_sha256": sha(encode(table)), + "protocol_sha256": "1" * 64, + "action_semantics": mixed._ACTION, + "pricing_semantics": mixed._PRICE_SEMANTICS, + "production_admission": True, + "source_and_measurements_verified": True, + "accuracy_pass": True, + "runtime_invariants_pass": True, + "matched_workload_pass": True, + "accuracy_result_sha256": "3" * 64, + "runtime_invariants_result_sha256": "4" * 64, + "matched_workload_result_sha256": "5" * 64, + "selector_replay_sha256": "6" * 64, + "measurement_results_sha256": sha(encode(["2" * 64])), + "economics": { + "policy_steps": 10, + "compact_choices": 4, + "gross_compact_value_ms_lower_bound": 3.0, + "fixed_confidence_path_overhead_ms_upper_bound": 0.1, + "safety_margin_ms": 0.5, + }, + } + return table, receipt, live + + +def resolve(table, receipt, live=None, *, update_receipt=True): + table_raw = encode(table) + receipt = copy.deepcopy(receipt) + if update_receipt: + receipt["costs_sha256"] = sha(table_raw) + receipt_raw = encode(receipt) + binding = mixed.MixedEvidenceBinding(sha(table_raw), sha(receipt_raw), "1" * 64) + return mixed.resolve_mixed_evidence( + binding, costs_raw=table_raw, admission_raw=receipt_raw, live_identity=live or identity() + ) + + +class TestMixedEvidence(unittest.TestCase): + def assert_native(self, table, receipt, live=None, **kwargs): + result = resolve(table, receipt, live, **kwargs) + self.assertIsNone(result.costs) + self.assertTrue(result.reason.startswith("mixed_evidence_invalid_native:")) + + def test_01_actual_binding_positive_synthetic_fixture_only(self): + table, receipt, live = fixture() + resolved = resolve(table, receipt, live) + self.assertEqual(resolved.reason, "mixed_evidence_bound") + self.assertEqual(resolved.costs.lookup(geometry()).candidate_upper_ms, 6.9) + self.assertEqual(resolved.costs.minimum_predicted_gain, 0.01) + + def test_02_bad_fixed_restored_old_validation_schema(self): + table, receipt, live = fixture() + table["schema"] = "dspark-host-ragged-bcg-prices-v1" + self.assert_native(table, receipt, live) + table["schema"] = "dspark-mixed-breakable-costs-v1" + self.assertIsNotNone(resolve(table, receipt, live).costs) + table["schema"] = "dspark-host-ragged-bcg-prices-v1" + self.assert_native(table, receipt, live) + + def test_03_absent_optional_binding_does_not_read_payloads(self): + result = mixed.resolve_mixed_evidence( + None, costs_raw=None, admission_raw=None, live_identity=identity() + ) + self.assertEqual(result.reason, "mixed_evidence_absent_native") + self.assertIsNone(result.costs) + + def test_04_decode_policy_object_and_flag_unchanged_for_all_outcomes(self): + decode = {"enabled": True, "identity": "existing-decode", "table": object()} + before = dict(decode) + table, receipt, live = fixture() + for binding in (None, mixed.MixedEvidenceBinding("0" * 64, "1" * 64, "2" * 64)): + mixed.resolve_mixed_evidence( + binding, costs_raw=b"bad", admission_raw=b"bad", live_identity=live + ) + self.assertEqual(decode, before) + self.assertTrue(decode["enabled"]) + self.assertIsNotNone(resolve(table, receipt, live).costs) + self.assertEqual(decode, before) + + def test_05_live_identity_is_not_read_from_table(self): + table, receipt, live = fixture() + for field in ( + "source_tree_sha256", + "runtime_snapshot_sha256", + "model_config_sha256", + "gpu_class_sha256", + "topology_sha256", + "selector_identity_sha256", + ): + with self.subTest(field=field): + self.assert_native(table, receipt, dataclasses.replace(live, **{field: "9" * 64})) + + def test_06_invalid_supported_runtime_and_mutable_grid_rejected(self): + live = identity() + for overrides in ( + {"physical_k": 2}, + {"physical_k": True}, + {"world_size": 1}, + {"execution_mode": "FULL"}, + {"body_buckets": list(live.body_buckets)}, + {"body_buckets": (96, 128, 192, 256)}, + ): + with self.subTest(overrides=overrides), self.assertRaises(ValueError): + dataclasses.replace(live, **overrides) + + def test_07_payloads_are_deeply_frozen_into_typed_tuples(self): + table, receipt, live = fixture() + costs = resolve(table, receipt, live).costs + table["cells"][0]["geometry"]["rank_features"][1]["eligible_generations"] = 0 + table["cells"][0]["native_lower_ms"] = 0 + self.assertEqual(costs.lookup(geometry()).native_lower_ms, 7.0) + with self.assertRaises(dataclasses.FrozenInstanceError): + costs.cells[0].candidate_upper_ms = 1 + + def test_08_uncovered_geometry_is_native_no_interpolation(self): + table, receipt, live = fixture() + costs = resolve(table, receipt, live).costs + unseen = dataclasses.replace(geometry(), actual_rank_rows=(64, 15, 0, 0, 0, 0, 0, 0)) + self.assertIsNone(costs.lookup(unseen)) + self.assertIsNone(costs.lookup(geometry().payload())) + + def test_09_duplicate_geometry_rejects_favorable_reselection(self): + table, receipt, _ = fixture() + table["cells"].append(copy.deepcopy(table["cells"][0])) + table["cells"][1]["candidate_upper_ms"] = 1 + self.assert_native(table, receipt) + + def test_10_cell_list_empty_or_unbounded_rejected(self): + table, receipt, _ = fixture() + for cells in ([], [table["cells"][0]] * 257): + bad = copy.deepcopy(table) + bad["cells"] = cells + self.assert_native(bad, receipt) + + def test_11_context_and_fixed_generation_rows_cannot_be_discounted(self): + original = geometry() + for change in ( + {"actual_rank_rows": (63, 16, 0, 0, 0, 0, 0, 0)}, + {"actual_rank_rows": (64, 7, 0, 0, 0, 0, 0, 0)}, + {"native_rank_rows": (64, 31, 0, 0, 0, 0, 0, 0)}, + ): + with self.subTest(change=change), self.assertRaises(ValueError): + dataclasses.replace(original, **change) + features = list(original.rank_features) + features[1] = (0, 0, 0, 0, 1024, 128, 1, 0, 7, 0, 8, 0) + with self.assertRaises(ValueError): + dataclasses.replace(original, rank_features=tuple(features)) + + def test_12_decode_only_and_noncompact_geometry_rejected(self): + original = geometry() + features = list(original.rank_features) + features[0] = (64, 4096, 8192, 0, 0, 0, 0, 0, 0, 1, 0, 0) + with self.assertRaises(ValueError): + dataclasses.replace(original, rank_features=tuple(features)) + with self.assertRaises(ValueError): + dataclasses.replace(original, actual_rank_rows=original.native_rank_rows) + + def test_13_rank_order_is_not_coalesced(self): + table, receipt, live = fixture() + costs = resolve(table, receipt, live).costs + original = geometry() + reordered = dataclasses.replace( + original, + rank_features=original.rank_features[::-1], + actual_rank_rows=original.actual_rank_rows[::-1], + native_rank_rows=original.native_rank_rows[::-1], + ) + self.assertIsNone(costs.lookup(reordered)) + + def test_14_unknown_null_and_validation_only_fields_reject(self): + table, receipt, _ = fixture() + for field, value in ( + ("validation_only", False), + ("pricing_sequence", [96]), + ("confidence_sps_table_path", "old.json"), + ("unknown", None), + ): + bad = copy.deepcopy(table) + bad[field] = value + self.assert_native(bad, receipt) + + def test_15_synthetic_prices_and_claim_only_receipts_reject(self): + table, receipt, _ = fixture() + for field in ( + "production_admission", + "source_and_measurements_verified", + "accuracy_pass", + "runtime_invariants_pass", + "matched_workload_pass", + ): + bad = copy.deepcopy(receipt) + bad[field] = False + self.assert_native(table, bad) + table["synthetic"] = True + self.assert_native(table, receipt) + + def test_16_proxy_is_never_relabelled_as_whole_step(self): + table, receipt, _ = fixture() + for key in ("table", "receipt"): + bad_table, bad_receipt = copy.deepcopy(table), copy.deepcopy(receipt) + (bad_table if key == "table" else bad_receipt)["pricing_semantics"] = ( + "qualified_group_whole_step" + ) + self.assert_native(bad_table, bad_receipt) + + def test_17_protocol_action_and_measurement_binding_reject_drift(self): + table, receipt, _ = fixture() + for key, value in ( + ("protocol_sha256", "9" * 64), + ("action_semantics", "uniform-L"), + ("measurement_results_sha256", "9" * 64), + ): + bad = copy.deepcopy(receipt) + bad[key] = value + self.assert_native(table, bad) + + def test_18_all_result_hashes_required(self): + table, receipt, _ = fixture() + for field in ( + "accuracy_result_sha256", + "runtime_invariants_result_sha256", + "matched_workload_result_sha256", + "selector_replay_sha256", + "measurement_results_sha256", + ): + bad = copy.deepcopy(receipt) + bad[field] = "not-a-hash" + self.assert_native(table, bad) + + def test_19_positive_economics_recomputed_no_boolean_override(self): + table, receipt, _ = fixture() + for change in ( + {"compact_choices": 0}, + {"gross_compact_value_ms_lower_bound": 1.0}, + {"policy_steps": True}, + {"safety_margin_ms": -1.0}, + {"compact_choices": 11}, + ): + bad = copy.deepcopy(receipt) + bad["economics"].update(change) + self.assert_native(table, bad) + + def test_20_no_relaxation_of_one_percent_immediate_gate(self): + table, receipt, _ = fixture() + for margin in (0.0, 0.009, True): + table["minimum_predicted_gain"] = margin + self.assert_native(table, receipt) + table["minimum_predicted_gain"] = 0.02 + self.assertEqual(resolve(table, receipt).costs.minimum_predicted_gain, 0.02) + + def test_21_nonpositive_and_nonfinite_costs_rejected(self): + table, receipt, _ = fixture() + for field in ("native_lower_ms", "candidate_upper_ms"): + for value in (0, -1, True, "7.0"): + bad = copy.deepcopy(table) + bad["cells"][0][field] = value + self.assert_native(bad, receipt) + raw = encode(table).replace(b'"candidate_upper_ms":6.9', b'"candidate_upper_ms":1e9999') + receipt["costs_sha256"] = sha(raw) + self.assert_raw_native(raw, encode(receipt)) + + def assert_raw_native(self, table_raw, receipt_raw): + binding = mixed.MixedEvidenceBinding(sha(table_raw), sha(receipt_raw), "1" * 64) + result = mixed.resolve_mixed_evidence( + binding, costs_raw=table_raw, admission_raw=receipt_raw, live_identity=identity() + ) + self.assertIsNone(result.costs) + + def test_22_duplicate_nonfinite_malformed_deep_json_rejected(self): + _, receipt, _ = fixture() + for raw in ( + b'{"schema":"x","schema":"y"}', + b'{"x":NaN}', + b'{"x":Infinity}', + b"{malformed", + b"\xff", + b"[" * 1500 + b"]" * 1500, + ): + self.assert_raw_native(raw, encode(receipt)) + + def test_23_missing_raw_oversized_and_digest_drift_rejected(self): + table, receipt, live = fixture() + raw, qraw = encode(table), encode(receipt) + binding = mixed.MixedEvidenceBinding(sha(raw), sha(qraw), "1" * 64) + for value in (None, raw + b" ", bytearray(raw), b"", b" " * (mixed._MAX_BYTES + 1)): + result = mixed.resolve_mixed_evidence( + binding, costs_raw=value, admission_raw=qraw, live_identity=live + ) + self.assertIsNone(result.costs) + + def test_24_table_receipt_cross_binding_rejected(self): + table, receipt, _ = fixture() + table["cells"][0]["candidate_upper_ms"] = 6.8 + self.assert_native(table, receipt, update_receipt=False) + + def test_25_rank_extents_strict_types_and_wire_limits(self): + original = geometry() + for value in (True, 1.0, -1, 2**63): + features = list(original.rank_features) + changed = list(features[0]) + changed[2] = value + features[0] = tuple(changed) + with self.subTest(value=value), self.assertRaises((TypeError, ValueError)): + dataclasses.replace(original, rank_features=tuple(features)) + + def test_26_decode_schema_is_not_mixed_evidence(self): + _, receipt, _ = fixture() + self.assert_raw_native(encode({"schema_version": 2}), encode(receipt)) + + def test_27_valid_bootstrap_reserves_full_k_plus_one_rows(self): + original = geometry() + features = list(original.rank_features) + features[1] = (0, 0, 0, 0, 1024, 128, 1, 0, 7, 0, 8, 4) + valid = dataclasses.replace( + original, rank_features=tuple(features), actual_rank_rows=(64, 20, 0, 0, 0, 0, 0, 0) + ) + self.assertEqual(valid.rank_features[1][-1], 4) + with self.assertRaises(ValueError): + dataclasses.replace(valid, actual_rank_rows=(64, 10, 0, 0, 0, 0, 0, 0)) + + def test_28_bad_fixed_restored_cost_pin_is_not_relabelled(self): + table, receipt, live = fixture() + raw, qraw = encode(table), encode(receipt) + bad_binding = mixed.MixedEvidenceBinding("0" * 64, sha(qraw), "1" * 64) + for binding, expected in ( + (bad_binding, False), + (mixed.MixedEvidenceBinding(sha(raw), sha(qraw), "1" * 64), True), + (bad_binding, False), + ): + result = mixed.resolve_mixed_evidence( + binding, costs_raw=raw, admission_raw=qraw, live_identity=live + ) + self.assertEqual(result.costs is not None, expected) + + def test_29_file_binding_requires_independent_live_fingerprint(self): + table, receipt, live = fixture() + table_raw, receipt_raw = encode(table), encode(receipt) + live_raw = encode( + {"schema": "dspark-mixed-live-fingerprint-v1", "identity": live.payload()} + ) + with tempfile.TemporaryDirectory() as directory: + paths = [Path(directory) / name for name in ("costs.json", "receipt.json", "live.json")] + for path, raw in zip(paths, (table_raw, receipt_raw, live_raw)): + path.write_bytes(raw) + config = SimpleNamespace( + costs_path=str(paths[0]), + costs_sha256=sha(table_raw), + admission_path=str(paths[1]), + admission_sha256=sha(receipt_raw), + protocol_sha256="1" * 64, + live_fingerprint_path=str(paths[2]), + live_fingerprint_sha256=sha(live_raw), + ) + self.assertIsNotNone(mixed.load_mixed_evidence(config).costs) + config.live_fingerprint_sha256 = "0" * 64 + self.assertIsNone(mixed.load_mixed_evidence(config).costs) + self.assertEqual(mixed.load_mixed_evidence(None).reason, "mixed_evidence_absent_native") + + def test_30_typed_config_rejects_invalid_pins_and_unknown_fields(self): + data = dict( + costs_path="costs.json", + costs_sha256="1" * 64, + admission_path="receipt.json", + admission_sha256="2" * 64, + protocol_sha256="3" * 64, + live_fingerprint_path="live.json", + live_fingerprint_sha256="4" * 64, + ) + self.assertEqual(DSparkMixedEvidenceConfig(**data).costs_path, "costs.json") + with self.assertRaises(ValueError): + DSparkMixedEvidenceConfig(**{**data, "costs_sha256": "invalid"}) + with self.assertRaises(ValueError): + DSparkMixedEvidenceConfig(**data, force=True) + + def test_31_feature_off_rejects_mixed_input_and_k1_clears_it(self): + config = SimpleNamespace( + enable_confidence_scheduling=False, + enable_fused_confidence_scheduler=False, + confidence_sts_path=None, + confidence_sps_table_path=None, + confidence_sps_live_fingerprint_path=None, + confidence_admission_receipt_path=None, + confidence_admission_receipt_sha256=None, + confidence_mixed_evidence=object(), + max_draft_len=3, + ) + with self.assertRaisesRegex(ValueError, "require enable_confidence_scheduling"): + DSparkDecodingConfig.validate_confidence_scheduling(config) + config.enable_confidence_scheduling, config.max_draft_len = True, 1 + self.assertIs(DSparkDecodingConfig.validate_confidence_scheduling(config), config) + self.assertFalse(config.enable_confidence_scheduling) + self.assertIsNone(config.confidence_mixed_evidence) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_runtime.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_runtime.py new file mode 100644 index 000000000000..73f8cb0aa618 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_runtime.py @@ -0,0 +1,350 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU contracts for optional DSpark mixed allocation; prices are synthetic.""" + +import copy +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor import dspark_mixed_runtime as runtime +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType +from tensorrt_llm._torch.speculative.dspark_mixed_evidence import ( + AdmittedMixedCosts, + MixedCostCell, + MixedEvidenceBinding, + MixedGeometry, + MixedRuntimeIdentity, +) +from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig +from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner + + +class _Event: + def __init__(self, ready=True): + self.ready = ready + self.records = 0 + + def query(self): + return self.ready + + def record(self): + self.records += 1 + + +def _request(request_id, *, context=False): + request = SimpleNamespace( + py_request_id=request_id, + py_seq_slot=request_id, + is_dummy=False, + py_batch_idx=0, + py_draft_tokens=[1, 2, 3], + py_verify_len=None, + state=SimpleNamespace(name="GENERATION_IN_PROGRESS"), + py_orig_prompt_len=128, + py_max_new_tokens=768, + context_chunk_size=64 if context else 0, + context_current_position=128 if context else 0, + py_num_compressed_tokens=0, + is_first_context_chunk=False, + ) + request.get_num_tokens = lambda beam: 128 + return request + + +def _group(monkeypatch, *, all_decode=False, stale=False): + monkeypatch.setattr(runtime, "_greedy", lambda _batch: True) + identity = MixedRuntimeIdentity(*(c * 64 for c in "abcdef")) + binding = MixedEvidenceBinding("1" * 64, "2" * 64, "3" * 64) + costs = AdmittedMixedCosts(identity, binding, 0.01, ()) + group = [] + for rank in range(8): + config = SimpleNamespace( + max_draft_len=3, + max_seq_len=4096, + max_num_tokens=4096, + enable_attention_dp=True, + torch_compile_enabled=False, + is_multimodal=False, + spec_config=SimpleNamespace( + max_draft_len=3, + block_size=3, + decoding_type="DSpark", + draft_is_embedded_in_target=True, + use_rejection_sampling=False, + ), + prefill_cuda_graph_backend=SimpleNamespace(name="BREAKABLE"), + prefill_cuda_graph_num_tokens=runtime.BODY_BUCKETS, + ) + generations = [_request(rank * 1000 + i + 1) for i in range(64)] + contexts = [_request(99999, context=True)] if rank == 0 and not all_decode else [] + batch = SimpleNamespace( + context_requests=contexts, + generation_requests=generations, + context_requests_last_chunk=contexts, + ) + planner = DSparkVerifyPlanner(cfg=DSparkScheduleConfig(3), device_windows=True) + planner._host_buffer = torch.full((64, 3), 5.0) + planner._host_confidence_stamp = torch.full((64,), 4, dtype=torch.int32) + planner._snapshot_valid, planner._copy_event = True, _Event() + slots = {r.py_request_id: i for i, r in enumerate(generations)} + incarnations = {r.py_request_id: i + 1 for i, r in enumerate(generations)} + worker = SimpleNamespace( + _draft_seq_host=5, + _req_to_slot=slots, + confidence_row_for=slots.__getitem__, + confidence_incarnation_for=incarnations.get, + ) + planner._mixed_current_snapshot_meta = dict( + buffer=planner._host_buffer, + event=planner._copy_event, + owners=slots.copy(), + staged_iteration=9, + staging_sequence=5, + confidence_stamp=planner._host_confidence_stamp, + producer_attempts={r.py_request_id: (r, i, i + 1) for i, r in enumerate(generations)}, + ) + if stale and rank == 1: + incarnations[generations[0].py_request_id] += 100 + runner = SimpleNamespace( + _config=config, + mapping=SimpleNamespace(tp_size=8, pp_size=1, cp_size=1), + _dspark_confidence_enabled=True, + _dspark_trims_submitted_tokens=True, + use_beam_search=False, + _dspark_device_budget=10, + cuda_graph_runner=SimpleNamespace( + agreed_ragged_bucket=192, + ragged_pad_verify_len=1, + ragged_zero_real_high_rows=0, + ), + breakable_cuda_graph_runner=SimpleNamespace( + is_capturing=False, is_warming_up=False, has_graph=lambda _bucket: True + ), + ) + group.append((runtime.MixedRuntime(runner, costs), batch, worker, planner, incarnations)) + votes = [ + mixed.local_vote(batch, worker, planner, 10, rank, executor_eligible=True) + for rank, (mixed, batch, worker, planner, _) in enumerate(group) + ] + if not all_decode: + index = runtime.BODY_BUCKETS.index(192) + geometry = MixedGeometry( + tuple(tuple(v[14:26]) for v in votes), + 192, + tuple(v[28 + index * 3 + 2] for v in votes), + tuple(v[27] for v in votes), + ) + costs = replace(costs, cells=(MixedCostCell(geometry, 10.0, 5.0, "4" * 64),)) + for mixed, *_ in group: + mixed.costs = costs + return group, votes + + +def test_host_mixed_shortens_decode_peer_without_mutating_context(monkeypatch): + group, votes = _group(monkeypatch) + context = copy.deepcopy(group[0][1].context_requests[0].__dict__) + for mixed, batch, worker, planner, _ in group: + assert mixed.apply(batch, worker, planner, votes, 10) + assert mixed.runner._dspark_host_window_step + assert mixed.runner.cuda_graph_runner._dspark_host_window_batch is batch + assert mixed.runner._dspark_device_budget is None + assert any(r.py_verify_len < 3 for r in batch.generation_requests) + assert group[0][1].context_requests[0].__dict__ == context + assert sum(r.py_verify_len + 1 for r in group[1][1].generation_requests) == 192 + + +@pytest.mark.parametrize("failure", ["uncovered", "unready", "identity"]) +def test_missing_mixed_evidence_stays_native_without_disabling_decode(monkeypatch, failure): + group, votes = _group(monkeypatch) + if failure == "uncovered": + for mixed, *_ in group: + mixed.costs = replace(mixed.costs, cells=()) + elif failure == "unready": + mixed, batch, worker, planner, _ = group[1] + planner._copy_event.ready = False + votes[1] = mixed.local_vote(batch, worker, planner, 10, 1, executor_eligible=True) + assert votes[1][5] == 0 + else: + votes[2][6] ^= 1 + for mixed, batch, worker, planner, _ in group: + assert mixed.apply(batch, worker, planner, votes, 10) + assert all(r.py_verify_len == 3 for r in batch.generation_requests) + assert mixed.runner._dspark_confidence_enabled + + +def test_stale_owner_keeps_full_width_without_discarding_other_owners(monkeypatch): + group, votes = _group(monkeypatch, stale=True) + mixed, batch, worker, planner, _ = group[1] + assert votes[1][22] == 63 + mixed.apply(batch, worker, planner, votes, 10) + assert batch.generation_requests[0].py_verify_len == 3 + assert any(r.py_verify_len < 3 for r in batch.generation_requests[1:]) + + +@pytest.mark.parametrize("change", ["incarnation", "context", "first_chunk"]) +def test_ticket_rechecked_before_publishing_lengths(monkeypatch, change): + group, votes = _group(monkeypatch) + rank = 1 if change == "incarnation" else 0 + mixed, batch, worker, planner, incarnations = group[rank] + if change == "incarnation": + incarnations[batch.generation_requests[0].py_request_id] += 1 + elif change == "context": + batch.context_requests[0].context_current_position += 1 + else: + batch.context_requests[0].is_first_context_chunk = True + with pytest.raises(ValueError, match="changed"): + mixed.apply(batch, worker, planner, votes, 10) + assert all(r.py_verify_len is None for r in batch.generation_requests) + + +@pytest.mark.parametrize("unresolved", ["first_chunk", "token_budget"]) +def test_unfinalized_geometry_declines_group_before_compact_publication(monkeypatch, unresolved): + group, votes = _group(monkeypatch) + mixed, batch, worker, planner, _ = group[0] + if unresolved == "first_chunk": + batch.context_requests[0].is_first_context_chunk = True + else: + mixed.runner._config.max_num_tokens = 319 + votes[0] = mixed.local_vote(batch, worker, planner, 10, 0, executor_eligible=True) + assert votes[0][5] == 0 + assert all(r.py_verify_len is None for r in batch.generation_requests) + for mixed, batch, worker, planner, _ in group: + assert mixed.apply(batch, worker, planner, votes, 10) + assert all(r.py_verify_len == 3 for r in batch.generation_requests) + assert mixed.runner._dspark_host_window_step + + +def test_first_chunk_stays_native_across_real_python_resource_preparation(monkeypatch): + group, votes = _group(monkeypatch) + mixed, batch, worker, planner, _ = group[0] + + class _NativeRequestModel(SimpleNamespace): + # Models the native first-chunk predicate; no C++ execution is claimed. + @property + def is_first_context_chunk(self): + return self.context_current_position == self.prepopulated_prompt_len + + values = batch.context_requests[0].__dict__.copy() + values.pop("is_first_context_chunk") + context = _NativeRequestModel(**values) + context.context_chunk_size, context.context_current_position = 256, 0 + context.prepopulated_prompt_len = 0 + batch.context_requests[:] = [context] + batch.context_requests_last_chunk[:] = [context] + votes[0] = mixed.local_vote(batch, worker, planner, 10, 0, executor_eligible=True) + assert votes[0][5] == 0 + for mixed, member, worker, planner, _ in group: + mixed.apply(member, worker, planner, votes, 10) + assert all(r.py_verify_len == 3 for r in member.generation_requests) + phases = [] + + def model_prefix_cache_setter(_batch): + phases.append("prepare") + context.prepopulated_prompt_len = context.context_current_position = 128 + context.context_chunk_size = 128 + + manager = SimpleNamespace( + resource_managers={ + ResourceManagerType.KV_CACHE_MANAGER: SimpleNamespace( + prepare_resources=model_prefix_cache_setter, + report_batch_to_connector=lambda _batch: phases.append("report"), + ) + }, + maybe_fit_token_budget=lambda _batch: phases.append("trim"), + ) + ResourceManager.prepare_resources(manager, batch) + assert phases == ["prepare", "trim", "report"] + assert context.context_chunk_size == 128 + assert all(r.py_verify_len == 3 for r in batch.generation_requests) + + +def test_all_decode_keeps_original_policy_and_device_mode(monkeypatch): + group, votes = _group(monkeypatch, all_decode=True) + for mixed, batch, worker, planner, _ in group: + before = [r.__dict__.copy() for r in batch.generation_requests] + assert not mixed.apply(batch, worker, planner, votes, 10) + assert [r.__dict__ for r in batch.generation_requests] == before + assert planner.device_windows + + +def test_optional_suffix_preserves_legacy_payload(monkeypatch): + group, votes = _group(monkeypatch) + mixed = group[0][0] + legacy, suffix = mixed.split_votes([[1, 2, 3] + vote for vote in votes], 3) + assert legacy == [[1, 2, 3]] * 8 + assert suffix == votes + with pytest.raises(ValueError, match="shape differs"): + mixed.split_votes([[1, 2, 3] + vote[:-1] for vote in votes], 3) + + +def test_unsupported_backend_declines_before_reading_snapshot(monkeypatch): + group, _ = _group(monkeypatch) + mixed, batch, worker, planner, _ = group[1] + mixed.runner._config.prefill_cuda_graph_backend = SimpleNamespace(name="PIECEWISE") + planner._mixed_current_snapshot_meta = None + vote = mixed.local_vote(batch, worker, planner, 10, 1, executor_eligible=True) + assert len(vote) == runtime.WIRE_WORDS + assert vote[5] == 0 + + +def test_host_metadata_skips_device_prologue_and_local_full_graph(): + runner = SimpleNamespace( + _dspark_trims_submitted_tokens=True, + _dspark_device_windows=True, + cuda_graph_runner=SimpleNamespace(agreed_ragged_bucket=192), + get_runtime_tokens_per_gen_step=lambda k: k + 1, + _ragged_token_lens=DecoderRunner._ragged_token_lens, + ) + metadata = SimpleNamespace( + runtime_tokens_per_gen_step=None, ragged_verify_lens=None, device_windows_mode=False + ) + requests = [SimpleNamespace(py_verify_len=1), SimpleNamespace(py_verify_len=3)] + DecoderRunner._publish_gen_token_layout( + runner, metadata, requests, runtime_draft_len=3, host_authoritative=True + ) + assert not metadata.device_windows_mode + assert metadata.ragged_verify_lens == [2, 4] + runner._dspark_host_window_step, runner._dspark_device_budget = True, 5 + assert not DecoderRunner._apply_device_window_prologue(runner, {}, None) + assert runner._dspark_device_budget is None + batch = SimpleNamespace(can_run_cuda_graph=True) + graph = SimpleNamespace( + _dspark_host_window_batch=batch, + is_encoder_decoder=False, + enable_encoder_decoder_mixed_cuda_graph=False, + ) + assert not CUDAGraphRunner._can_run_cuda_graph_batch(graph, batch) + + +def test_producer_stamp_copy_uses_same_event_and_validates_before_rotation(monkeypatch): + monkeypatch.setattr(torch.cuda, "Event", _Event) + planner = DSparkVerifyPlanner(cfg=DSparkScheduleConfig(3), device_windows=True) + planner.stage_confidence(torch.ones(2, 3), torch.tensor([1, 2], dtype=torch.int32)) + first = (planner._host_buffer, planner._host_confidence_stamp, planner._copy_event) + with pytest.raises(ValueError, match="one per slot"): + planner.stage_confidence(torch.zeros(2, 3), torch.ones(3, dtype=torch.int32)) + assert planner._host_buffer is first[0] + assert planner._host_confidence_stamp is first[1] + planner.stage_confidence(torch.zeros(2, 3), torch.tensor([3, 4], dtype=torch.int32)) + assert planner._prev_buffer is first[0] + assert planner._prev_confidence_stamp is first[1] + assert planner._prev_event is first[2] + assert planner._host_confidence_stamp.tolist() == [3, 4] + + +def test_forward_attempt_alone_is_not_confidence_publication(monkeypatch): + group, _ = _group(monkeypatch) + _, batch, worker, planner, _ = group[1] + runtime.stage_mixed_ownership(planner, worker, 10, batch.generation_requests) + assert planner._mixed_current_snapshot_meta["producer_attempts"] == {} + worker._draft_seq_host += 1 + runtime.stage_mixed_ownership(planner, worker, 11, batch.generation_requests) + assert len(planner._mixed_current_snapshot_meta["producer_attempts"]) == 64 + assert ( + planner._mixed_current_snapshot_meta["confidence_stamp"] is planner._host_confidence_stamp + ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py new file mode 100644 index 000000000000..b1669101c24c --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py @@ -0,0 +1,346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Focused CPU tests for DSpark decoder-runner ragged runtime state.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine + + +def _runner(*, secondary=True): + return SimpleNamespace( + enabled=True, + agreed_ragged_bucket=None, + ragged_pad_verify_len=0, + ragged_zero_real_high_rows=0, + supported_batch_sizes=[4, 16], + secondary_padding_dummy_requests={5: object()} if secondary else {}, + _round_up_batch_size=lambda rows: 4 if rows <= 4 else 16, + will_pad_to=lambda *_args: True, + ) + + +def _engine(runner, buckets): + engine = object.__new__(DecoderRunner) + engine.cuda_graph_runner = runner + engine._config = SimpleNamespace( + spec_config=SimpleNamespace(max_draft_len=5), + max_batch_size=8, + cuda_graph_batch_sizes=[1, 4, 8], + use_mrope=False, + ) + engine._dspark_confidence_enabled = True + engine._dspark_trims_submitted_tokens = True + engine._dspark_device_windows = False + engine._dspark_last_padded_bs = None + engine.ragged_verify_token_buckets = lambda _rows: list(buckets) + return engine + + +def test_pinned_host_ping_pongs_without_private_cuda_events() -> None: + engine = SimpleNamespace( + _pinned_host_cache={}, + _pinned_host_active={}, + ) + + first = DecoderRunner._pinned_host(engine, "rows", [1, 2], torch.long) + second = DecoderRunner._pinned_host(engine, "rows", [3, 4], torch.long) + third = DecoderRunner._pinned_host(engine, "rows", [5, 6], torch.long) + + assert first.data_ptr() == third.data_ptr() + assert first.data_ptr() != second.data_ptr() + assert third.tolist() == [5, 6] + + +@pytest.mark.parametrize( + ("verifier_budget", "scheduled_window", "high_rows"), + [(80, 4, 0), (81, 5, 1), (88, 5, 8)], +) +def test_zero_real_exact_fit_uses_one_scheduled_dummy(verifier_budget, scheduled_window, high_rows): + runner = _runner() + engine = _engine(runner, [48, 64, 80, 81, 88, 96]) + dummy = SimpleNamespace(is_attention_dp_dummy=True, py_verify_len=None) + + bucket = DecoderRunner.fit_ragged_verify_lens( + engine, + [dummy], + [scheduled_window], + peer_stats=[[16, 88, 1], [0, 0, 1]], + exact_shape=(16, verifier_budget, 5), + exact_zero_real=True, + ) + + assert bucket == verifier_budget + assert dummy.py_verify_len == scheduled_window + assert engine._dspark_last_num_real == 0 + assert runner.ragged_pad_verify_len == 4 + assert runner.ragged_zero_real_high_rows == high_rows + + +def test_zero_real_exact_fit_requires_the_secondary_dummy(): + runner = _runner(secondary=False) + engine = _engine(runner, [88]) + dummy = SimpleNamespace(is_attention_dp_dummy=True, py_verify_len=None) + + with pytest.raises(RuntimeError, match="disappeared after"): + DecoderRunner.fit_ragged_verify_lens( + engine, + [dummy], + [5], + peer_stats=[[16, 88, 1], [0, 0, 1]], + exact_shape=(16, 88, 5), + exact_zero_real=True, + ) + + +def test_exact_fit_publishes_the_measured_bucket_and_pad_window(): + runner = _runner() + engine = _engine(runner, [8, 24]) + requests = [SimpleNamespace(py_verify_len=None) for _ in range(2)] + + bucket = DecoderRunner.fit_ragged_verify_lens( + engine, + requests, + [2, 2], + peer_stats=[[2, 6, 1]], + exact_shape=(4, 8, 1), + ) + + assert bucket == 8 + assert [request.py_verify_len for request in requests] == [2, 2] + assert runner.agreed_ragged_bucket == 8 + assert runner.ragged_pad_verify_len == 0 + + +def test_full_k_bucket_preserves_the_native_static_graph(): + runner = _runner() + engine = _engine(runner, [24]) + requests = [SimpleNamespace(py_verify_len=None) for _ in range(2)] + + bucket = DecoderRunner.fit_ragged_verify_lens( + engine, + requests, + [5, 5], + peer_stats=[[2, 12, 1]], + exact_shape=(4, 24, 6), + ) + + assert bucket is None + assert runner.agreed_ragged_bucket is None + assert all(request.py_verify_len is None for request in requests) + + +def _warmup_metadata(position_helper, mask_helper): + ratios = [1, 4] + return SimpleNamespace( + _compress_ratios_sorted=ratios, + max_draft_tokens=5, + past_kv_lens_cuda={ratio: torch.empty(8, dtype=torch.int32) for ratio in ratios}, + cu_new_comp_kv_cuda={ratio: torch.empty(9, dtype=torch.int32) for ratio in ratios}, + new_comp_kv_lens_cuda={ratio: torch.empty(8, dtype=torch.int32) for ratio in ratios}, + compressed_position_ids_cuda={ + ratio: torch.empty(48, dtype=torch.int32) for ratio in ratios + }, + compressed_mask_cuda={ratio: torch.empty(48, dtype=torch.bool) for ratio in ratios}, + _compute_gen_compressed_position_ids=position_helper, + _compute_compressed_mask=mask_helper, + ) + + +def test_ragged_compressor_warmup_primes_and_cleans_intermediate_shape(monkeypatch): + position_calls = [] + mask_calls = [] + + def record_positions(*args): + position_calls.append((torch.is_inference_mode_enabled(), args[3:])) + + def record_mask(*args): + mask_calls.append((torch.is_inference_mode_enabled(), args[3:])) + + metadata = _warmup_metadata(record_positions, record_mask) + engine = SimpleNamespace( + _dspark_confidence_enabled=True, + _dspark_trims_submitted_tokens=True, + attn_metadata=metadata, + _config=SimpleNamespace(max_batch_size=8, cuda_graph_batch_sizes=[1, 4, 8]), + ) + sync_calls = [] + monkeypatch.setattr(torch.cuda, "synchronize", lambda: sync_calls.append(True)) + + DecoderRunner._warmup_dspark_ragged_compressor_metadata(engine) + + assert position_calls == [(True, (0, 7, 6, [1, 4], {1: 0, 4: 0}))] + assert mask_calls == [(True, (7, {1: 42, 4: 14}, [1, 4]))] + for ratio, token_count in ((1, 42), (4, 14)): + assert not metadata.past_kv_lens_cuda[ratio][:7].any() + assert not metadata.new_comp_kv_lens_cuda[ratio][:7].any() + assert not metadata.cu_new_comp_kv_cuda[ratio][:8].any() + assert not metadata.compressed_position_ids_cuda[ratio][:token_count].any() + assert not metadata.compressed_mask_cuda[ratio][:token_count].any() + assert sync_calls == [True, True] + + +def test_ragged_compressor_warmup_cleans_after_helper_failure(monkeypatch): + def fail_positions(*args): + args[2][1][:18].fill_(7) + raise RuntimeError("compile boom") + + metadata = _warmup_metadata(fail_positions, lambda *_args: None) + engine = SimpleNamespace( + _dspark_confidence_enabled=True, + _dspark_trims_submitted_tokens=True, + attn_metadata=metadata, + _config=SimpleNamespace(max_batch_size=4, cuda_graph_batch_sizes=[1, 4]), + ) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + + with pytest.raises(RuntimeError, match="compile boom"): + DecoderRunner._warmup_dspark_ragged_compressor_metadata(engine) + + assert not metadata.past_kv_lens_cuda[1][:3].any() + assert not metadata.new_comp_kv_lens_cuda[1][:3].any() + assert not metadata.cu_new_comp_kv_cuda[1][:4].any() + assert not metadata.compressed_position_ids_cuda[1][:18].any() + assert not metadata.compressed_mask_cuda[1][:18].any() + + +def test_ragged_compressor_warmup_requires_trimmed_tokens(): + engine = SimpleNamespace( + _dspark_confidence_enabled=True, + _dspark_trims_submitted_tokens=False, + ) + + DecoderRunner._warmup_dspark_ragged_compressor_metadata(engine) + + +def test_ragged_compressor_warmup_rejects_missing_metadata_contract(): + engine = SimpleNamespace( + _dspark_confidence_enabled=True, + _dspark_trims_submitted_tokens=True, + attn_metadata=SimpleNamespace(), + ) + + with pytest.raises(RuntimeError, match="missing .*max_draft_tokens"): + DecoderRunner._warmup_dspark_ragged_compressor_metadata(engine) + + +@pytest.mark.parametrize("runtime_draft_len", [0, 2, 5]) +def test_publish_gen_token_layout_uses_call_draft_length(runtime_draft_len: int) -> None: + engine = _engine(_runner(), []) + engine.get_runtime_tokens_per_gen_step = lambda draft_len: draft_len + 1 + requests = [ + SimpleNamespace(py_verify_len=min(1, runtime_draft_len)), + SimpleNamespace(py_verify_len=min(3, runtime_draft_len)), + ] + metadata = SimpleNamespace( + runtime_tokens_per_gen_step=0, + ragged_verify_lens=None, + device_windows_mode=False, + ) + assert not hasattr(engine, "runtime_draft_len") + + engine._publish_gen_token_layout(metadata, requests, runtime_draft_len=runtime_draft_len) + + assert metadata.runtime_tokens_per_gen_step == runtime_draft_len + 1 + assert metadata.ragged_verify_lens == [request.py_verify_len + 1 for request in requests] + assert not metadata.device_windows_mode + assert not hasattr(engine, "runtime_draft_len") + + +def test_attach_ragged_layout_refreshes_persistent_buffers_in_place() -> None: + runner = _engine(_runner(), []) + runner._pinned_host_cache = {} + runner._pinned_host_active = {} + runner.ragged_verify_lens_cuda = torch.empty(9, dtype=torch.int32) + runner.ragged_qo_indptr_cuda = torch.empty(10, dtype=torch.int32) + attention = SimpleNamespace(ragged_verify_lens=None) + speculation = SimpleNamespace() + requests = [SimpleNamespace(py_verify_len=1), SimpleNamespace(py_verify_len=3)] + + runner._attach_ragged_verify_layout(speculation, attention, requests) + first_lens_ptr = speculation.verify_lens.data_ptr() + first_prefix_ptr = speculation.qo_indptr.data_ptr() + + assert speculation.verify_lens.tolist() == [2, 4] + assert speculation.qo_indptr.tolist() == [0, 2, 6] + assert speculation.total_verify_tokens == 6 + assert attention.ragged_verify_lens == [2, 4] + + requests[0].py_verify_len = 0 + requests[1].py_verify_len = 2 + runner._attach_ragged_verify_layout(speculation, attention, requests) + + assert speculation.verify_lens.tolist() == [1, 3] + assert speculation.qo_indptr.tolist() == [0, 1, 4] + assert speculation.total_verify_tokens == 4 + assert attention.ragged_verify_lens == [1, 3] + assert speculation.verify_lens.data_ptr() == first_lens_ptr + assert speculation.qo_indptr.data_ptr() == first_prefix_ptr + + +@pytest.mark.parametrize("trim_enabled", [False, True]) +def test_missing_or_disabled_windows_clear_stale_ragged_views(trim_enabled: bool) -> None: + runner = _engine(_runner(), []) + runner._dspark_trims_submitted_tokens = trim_enabled + attention = SimpleNamespace(ragged_verify_lens=[2, 4]) + speculation = SimpleNamespace( + verify_lens=torch.tensor([2, 4]), + qo_indptr=torch.tensor([0, 2, 6]), + total_verify_tokens=6, + ) + + runner._attach_ragged_verify_layout( + speculation, + attention, + [SimpleNamespace(py_verify_len=1), SimpleNamespace(py_verify_len=None)], + ) + + assert speculation.verify_lens is None + assert speculation.qo_indptr is None + assert speculation.total_verify_tokens is None + assert attention.ragged_verify_lens is None + + +def test_model_engine_delegates_ragged_state_to_decoder_owner() -> None: + runner = _engine(_runner(), []) + runner._dspark_sps_cost_table = object() + runner._dspark_exact_candidate_cells = ((4, 12),) + runner._dspark_exact_identity_words = tuple(range(8)) + runner.fit_ragged_verify_lens = lambda *args, **kwargs: (args, kwargs) + worker = object() + runner._get_spec_worker = lambda: worker + engine = object.__new__(PyTorchModelEngine) + engine._cleanup_done = True + engine._runner = runner + + assert engine._dspark_confidence_enabled + assert engine._dspark_trims_submitted_tokens + assert engine._dspark_sps_cost_table is runner._dspark_sps_cost_table + assert engine._dspark_exact_candidate_cells == ((4, 12),) + assert engine._dspark_exact_identity_words == tuple(range(8)) + assert engine._dspark_device_budget is None + assert engine._get_spec_worker() is worker + engine._dspark_device_budget = 7 + assert runner._dspark_device_budget == 7 + assert engine.fit_ragged_verify_lens([1], exact_shape=(4, 12, 3)) == ( + ([1],), + {"exact_shape": (4, 12, 3)}, + ) + + +def test_model_engine_non_decoder_keeps_confidence_inactive() -> None: + engine = object.__new__(PyTorchModelEngine) + engine._cleanup_done = True + engine._runner = SimpleNamespace() + + assert not engine._dspark_confidence_enabled + assert not engine._dspark_trims_submitted_tokens + assert engine._dspark_sps_cost_table is None + assert engine._dspark_exact_candidate_cells == () + assert engine._dspark_exact_identity_words == (0,) * 8 + assert engine._dspark_device_budget is None + assert engine._get_spec_worker() is None diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sampling_admission.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sampling_admission.py new file mode 100644 index 000000000000..962cf4e5a7b6 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sampling_admission.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unsupported confidence/sampler combinations fail before model execution.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm._torch.speculative.dspark import validate_confidence_request +from tensorrt_llm._torch.speculative.interface import SpecMetadata + +pytestmark = pytest.mark.cpu_only + + +@pytest.mark.parametrize("enabled,rejection", [(False, True), (True, False), (True, True)]) +@pytest.mark.parametrize( + "controls,greedy", + [ + ({}, True), + ({"temperature": 0.0}, True), + ({"temperature": 1.0}, False), + ({"temperature": 1.0, "top_k": 1}, True), + ({"min_p": 0.2}, False), + ({"min_p": 1.0}, True), + ], +) +def test_request_admission_matches_actual_draft_sampling_controls( + enabled, rejection, controls, greedy +): + config = SimpleNamespace(enable_confidence_scheduling=enabled, use_rejection_sampling=rejection) + sampling = SimpleNamespace(temperature=None, top_k=None, top_p=None, min_p=None) + sampling.__dict__.update(controls) + request = SimpleNamespace(sampling_config=sampling, py_logits_post_processors=None) + if enabled and rejection and not greedy: + with pytest.raises(ValueError, match="advanced rejection draft sampling"): + validate_confidence_request(config, request) + else: + validate_confidence_request(config, request) + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_per_request_logit_processor_is_rejected_at_confidence_admission(enabled): + config = SimpleNamespace(enable_confidence_scheduling=enabled, use_rejection_sampling=False) + request = SimpleNamespace(sampling_config=None, py_logits_post_processors=[object()]) + if enabled: + with pytest.raises(ValueError, match="per-request logits post-processors"): + validate_confidence_request(config, request) + else: + validate_confidence_request(config, request) + + +@pytest.mark.parametrize( + "enabled,rejection,flags", + [ + (True, True, [True, True]), + (True, True, [True, False]), + (True, True, [False, True]), + (True, True, [False]), + (True, False, [False, True]), + (False, True, [False, True]), + ], +) +def test_host_sampler_guard_is_group_uniform_before_graph_lookup(enabled, rejection, flags): + errors = [] + for local_flag in flags: + runner = object.__new__(DecoderRunner) + runner._dspark_confidence_enabled = enabled + runner.mapping = SimpleNamespace(enable_lm_head_tp_in_adp=False, tp_size=len(flags)) + exchange = Mock(return_value=torch.tensor([[flag] for flag in flags], dtype=torch.int64)) + runner.dist = SimpleNamespace(tp_allgather_int64=exchange) + metadata = object.__new__(SpecMetadata) + metadata.use_rejection_sampling = rejection + metadata.is_all_greedy_sample = local_flag + metadata.group_all_greedy_sample = None + runner._sync_group_all_greedy_sample(metadata) + if enabled and rejection and not all(flags): + with pytest.raises(ValueError, match="advanced rejection draft sampling"): + runner._validate_dspark_draft_sampling(metadata) + errors.append(True) + else: + runner._validate_dspark_draft_sampling(metadata) + errors.append(False) + assert exchange.call_count == int(enabled and rejection and len(flags) > 1) + assert len(set(errors)) == 1 + + +@pytest.mark.parametrize( + "enabled,rejection", [(False, False), (False, True), (True, False), (True, True)] +) +def test_capture_gate_retains_only_supported_advanced_variants(enabled, rejection): + runner = object.__new__(DecoderRunner) + runner._dspark_confidence_enabled = enabled + runner._config = SimpleNamespace(spec_config=SimpleNamespace(use_rejection_sampling=rejection)) + assert runner._supports_advanced_draft_sampling() is not (enabled and rejection) + + +@pytest.mark.parametrize("processors", [False, True]) +def test_invalid_confidence_request_is_excluded_before_activation(processors): + executor = object.__new__(PyExecutor) + executor.model_engine = SimpleNamespace( + _dspark_confidence_enabled=True, + spec_config=SimpleNamespace(enable_confidence_scheduling=True, use_rejection_sampling=True), + ) + request = SimpleNamespace( + sampling_config=SimpleNamespace(temperature=1.0, top_k=None, top_p=None, min_p=None), + py_logits_post_processors=[object()] if processors else None, + ) + already_active = object() + executor.waiting_queue = object() + executor.active_requests = [already_active] + executor._fetch_new_requests = Mock(return_value=[request]) + executor._handle_errors = Mock() + executor._configure_csa2_remote_tail = Mock() + + assert executor._fetch_and_activate_new_requests() == [] + assert executor.active_requests == [already_active] + executor._configure_csa2_remote_tail.assert_not_called() + executor._handle_errors.assert_called_once() + assert executor._handle_errors.call_args.kwargs == { + "requests": [request], + "charge_budget": False, + } + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_feature_off_retains_existing_lm_head_tp_group_sync(enabled): + runner = object.__new__(DecoderRunner) + runner._dspark_confidence_enabled = enabled + runner.mapping = SimpleNamespace(enable_lm_head_tp_in_adp=True, tp_size=2) + exchange = Mock(return_value=torch.tensor([[True], [False]], dtype=torch.int64)) + runner.dist = SimpleNamespace(tp_allgather_int64=exchange) + metadata = object.__new__(SpecMetadata) + metadata.use_rejection_sampling = True + metadata.is_all_greedy_sample = True + metadata.group_all_greedy_sample = None + + runner._sync_group_all_greedy_sample(metadata) + + exchange.assert_called_once_with([True]) + assert not metadata.group_all_greedy_sample + assert not metadata.is_all_greedy_sample diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_schedule_fused.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_schedule_fused.py new file mode 100644 index 000000000000..da592b60072a --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_schedule_fused.py @@ -0,0 +1,543 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm._torch.speculative.dspark_device_select as device_select +import tensorrt_llm._torch.speculative.dspark_schedule as dspark_schedule +from tensorrt_llm._torch.speculative.dspark_device_select import select_windows_device +from tensorrt_llm._torch.speculative.dspark_schedule import ( + DSparkScheduleConfig, + schedule_verify_lens_topk, + schedule_verify_lens_topk_fused_fill, +) +from tensorrt_llm._torch.speculative.ragged_helpers import fill_bucket_device + + +def _legacy_schedule_and_fill( + *, + survival: torch.Tensor, + budget: int, + num_real: int, + pad_len: int, + graph_num_tokens: int, + cfg: DSparkScheduleConfig, +) -> torch.Tensor: + rows = torch.arange(survival.shape[0], device=survival.device) + real_survival = torch.where( + (rows < num_real).unsqueeze(1), + survival, + torch.zeros_like(survival), + ) + scheduled = schedule_verify_lens_topk(survival=real_survival, budget=budget, cfg=cfg) + return fill_bucket_device( + scheduled + 1, + num_real=torch.tensor(num_real, device=survival.device), + graph_num_tokens=graph_num_tokens, + max_verify_len=cfg.resolved_max_verify_len + 1, + pad_fill=pad_len, + ) + + +@pytest.mark.parametrize( + "num_rows,num_real,budget,survival_eps", + [ + (1, 0, 0, 1e-6), + (1, 1, 0, 1e-6), + (2, 1, 1, 1e-6), + (2, 1, 9, 1e-6), + (4, 2, 3, 1e-6), + (4, 4, 16, 1e-6), + (8, 4, 9, 1e-6), + (2, 2, 1, 0.35), + (4, 1, 3, 0.35), + (8, 8, 32, 0.35), + ], +) +def test_cpu_fused_helper_matches_established_schedule_and_fill( + num_rows, num_real, budget, survival_eps +): + cfg = DSparkScheduleConfig( + block_size=5, + min_verify_len=1, + max_verify_len=5, + survival_eps=survival_eps, + ) + generator = torch.Generator().manual_seed( + 10000 * num_rows + 1000 * num_real + 10 * budget + int(survival_eps * 10) + ) + survival = torch.cumprod(torch.rand(num_rows, 5, generator=generator), dim=1) + survival[num_real:] = 0 + pad_len = 1 + token_floor = cfg.min_verify_len + 1 + max_token_len = cfg.resolved_max_verify_len + 1 + pad_tokens = (num_rows - num_real) * pad_len + capacity = min(max(budget, 0), num_real * cfg.schedulable_per_request) + minimum = num_real * token_floor + pad_tokens + maximum = num_real * max_token_len + pad_tokens + graph_num_tokens = min(maximum, minimum + capacity + min(num_real, 2)) + + expected = _legacy_schedule_and_fill( + survival=survival, + budget=budget, + num_real=num_real, + pad_len=pad_len, + graph_num_tokens=graph_num_tokens, + cfg=cfg, + ) + actual = schedule_verify_lens_topk_fused_fill( + survival=survival, + budget=budget, + num_real=num_real, + pad_len=pad_len, + graph_num_tokens=graph_num_tokens, + cfg=cfg, + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("num_real,budget,graph_tokens", [(2, 2, 9), (4, 4, 14)]) +def test_fused_workspace_selection_matches_tensor_and_reuses_storage( + num_real, budget, graph_tokens +): + cfg = DSparkScheduleConfig(block_size=3, min_verify_len=1) + workspace = device_select.DeviceWindowWorkspace.allocate( + max_rows=6, max_tokens=16, device="cpu" + ) + workspace.verify_lens.fill_(-777) + pointers = { + name: getattr(workspace, name).data_ptr() + for name in ("verify_lens", "qo_indptr", "req_idx", "kv_correction") + } + logits = torch.tensor([[10.0, 5.0, 1.0], [-2.0, -3.0, -4.0], [0.5, 0.2, 0.1], [1.0, 2.0, 3.0]]) + for confidence in (logits, logits.flip(0)): + kwargs = dict( + confidence_logits=confidence, + slot_idx=torch.tensor([0, 1, 2, 3]), + num_real=num_real, + budget=budget, + graph_num_tokens=graph_tokens, + cfg=cfg, + pad_len=1, + ) + expected = select_windows_device(**kwargs, use_fused_exact=False) + actual = select_windows_device(**kwargs, use_fused_exact=True, workspace=workspace) + assert expected.workspace is None + assert actual.workspace is workspace + for name, pointer in pointers.items(): + result = getattr(actual, name) + assert result.data_ptr() == pointer + torch.testing.assert_close(result, getattr(expected, name), rtol=0, atol=0) + assert expected.verify_lens.data_ptr() != workspace.verify_lens.data_ptr() + assert workspace.verify_lens[4:].tolist() == [-777, -777] + + +@pytest.mark.parametrize("num_rows", [0, 2]) +def test_fused_helper_preserves_output_address_and_unused_capacity(num_rows): + cfg = DSparkScheduleConfig(block_size=3, min_verify_len=1) + output = torch.full((num_rows + 2,), -777, dtype=torch.int32) + actual = schedule_verify_lens_topk_fused_fill( + survival=torch.ones(num_rows, 3), + budget=0, + num_real=num_rows, + pad_len=1, + graph_num_tokens=2 * num_rows, + cfg=cfg, + out=output, + ) + assert actual.untyped_storage().data_ptr() == output.untyped_storage().data_ptr() + assert actual.tolist() == [2] * num_rows + assert output[num_rows:].tolist() == [-777, -777] + + +@pytest.mark.parametrize( + "output,error,message", + [ + (torch.empty(2, dtype=torch.int32, device="meta"), ValueError, "share survival.device"), + (torch.empty(2, dtype=torch.int64), TypeError, "dtype torch.int32"), + (torch.empty(1, 2, dtype=torch.int32), ValueError, "1-D tensor"), + (torch.empty(1, dtype=torch.int32), ValueError, "at least 2 elements"), + (torch.empty(4, dtype=torch.int32)[::2], ValueError, "contiguous"), + ], +) +def test_fused_helper_rejects_incompatible_output_before_dispatch(output, error, message): + cfg = DSparkScheduleConfig(block_size=3, min_verify_len=1) + with pytest.raises(error, match=message): + schedule_verify_lens_topk_fused_fill( + survival=torch.ones(2, 3), + budget=0, + num_real=2, + pad_len=1, + graph_num_tokens=4, + cfg=cfg, + out=output, + ) + + +def test_round_robin_slack_is_not_promoted_into_confidence_budget(): + cfg = DSparkScheduleConfig( + block_size=5, + min_verify_len=1, + max_verify_len=5, + survival_eps=0.5, + ) + survival = torch.tensor( + [ + [1.0, 0.9, 0.1, 0.1, 0.1], + [1.0, 0.8, 0.49, 0.1, 0.1], + ] + ) + actual = schedule_verify_lens_topk_fused_fill( + survival=survival, + budget=2, + num_real=2, + pad_len=1, + graph_num_tokens=7, + cfg=cfg, + ) + # Budget two selects one position per row, then the already-paid graph + # remainder goes to row zero. Promoting that remainder into top-k would + # instead select row one's 0.49 candidate and produce [3, 4]. + assert actual.tolist() == [4, 3] + + +def test_zero_epsilon_declines_fusion_and_preserves_tensor_semantics(monkeypatch): + cfg = DSparkScheduleConfig( + block_size=5, + min_verify_len=1, + max_verify_len=5, + survival_eps=0.0, + ) + + def _must_not_run(**_kwargs): + raise AssertionError("zero-epsilon selector must use the tensor path") + + monkeypatch.setattr(device_select, "schedule_verify_lens_topk_fused_fill", _must_not_run) + result = select_windows_device( + confidence_logits=torch.zeros(3, 5), + slot_idx=torch.tensor([0, 1, 2]), + num_real=2, + budget=2, + graph_num_tokens=7, + cfg=cfg, + pad_len=1, + use_fused_exact=True, + ) + assert int(result.verify_lens.sum()) == 7 + + +def test_tensor_controls_decline_fusion(monkeypatch): + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1, max_verify_len=5) + + def _must_not_run(**_kwargs): + raise AssertionError("tensor controls must use the capture-safe tensor path") + + monkeypatch.setattr(device_select, "schedule_verify_lens_topk_fused_fill", _must_not_run) + result = select_windows_device( + confidence_logits=torch.zeros(2, 5), + slot_idx=torch.tensor([0, 1]), + num_real=torch.tensor(2, dtype=torch.int64), + budget=torch.tensor(1, dtype=torch.int64), + graph_num_tokens=5, + cfg=cfg, + pad_len=1, + use_fused_exact=True, + ) + assert int(result.verify_lens.sum()) == 5 + + +def test_fused_helper_rejects_infeasible_graph_without_dispatch(): + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1, max_verify_len=5) + with pytest.raises(ValueError, match="cannot realize"): + schedule_verify_lens_topk_fused_fill( + survival=torch.ones(2, 5), + budget=4, + num_real=2, + pad_len=1, + graph_num_tokens=5, + cfg=cfg, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("graph_bs", [16, 32, 64, 128, 192, 256]) +def test_cuda_fused_matches_tensor_oracle_across_all_graph_sizes(graph_bs): + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1, max_verify_len=5) + generator = torch.Generator(device="cuda").manual_seed(347805 + graph_bs) + confidence = torch.rand((graph_bs, 5), generator=generator, device="cuda") + survival = torch.cumprod(confidence, dim=1) + # Include exact ties and non-finite candidates in every graph shape. + survival[:2] = 0.75 + survival[2, 2] = torch.nan + survival[3, 1] = torch.inf + num_real = graph_bs - 3 + survival[num_real:] = 0 + budget = min(2 * num_real, num_real * cfg.schedulable_per_request) + pad_len = 1 + minimum = num_real * (cfg.min_verify_len + 1) + (graph_bs - num_real) + graph_num_tokens = minimum + budget + num_real + + expected = _legacy_schedule_and_fill( + survival=survival, + budget=budget, + num_real=num_real, + pad_len=pad_len, + graph_num_tokens=graph_num_tokens, + cfg=cfg, + ) + actual = schedule_verify_lens_topk_fused_fill( + survival=survival, + budget=budget, + num_real=num_real, + pad_len=pad_len, + graph_num_tokens=graph_num_tokens, + cfg=cfg, + ) + torch.cuda.synchronize() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_model_engine_prewarm_requires_device_windows(): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + + engine = object.__new__(DecoderRunner) + engine._dspark_fused_scheduler_enabled = True + engine._dspark_confidence_enabled = True + engine._dspark_trims_submitted_tokens = True + engine._dspark_device_windows = False + + def _must_not_get_worker(): + raise AssertionError("disabled device-window mode must not prewarm fusion") + + engine._get_spec_worker = _must_not_get_worker + DecoderRunner._warmup_dspark_fused_scheduler(engine) + assert engine._dspark_fused_schedule_ready_sizes == set() + + +def test_model_engine_prewarm_records_only_successful_graph_sizes(monkeypatch): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import runner as decoder_runner + + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1, max_verify_len=5) + planner = SimpleNamespace( + cfg=cfg, + exact_cost_table=SimpleNamespace(tables={16: object(), 32: object()}), + ) + engine = object.__new__(DecoderRunner) + engine._dspark_fused_scheduler_enabled = True + engine._dspark_confidence_enabled = True + engine._dspark_trims_submitted_tokens = True + engine._dspark_device_windows = True + engine._get_spec_worker = lambda: SimpleNamespace(verify_planner=planner) + + original_ones = torch.ones + monkeypatch.setattr( + decoder_runner.torch, + "ones", + lambda shape, dtype=None, device=None: original_ones(shape, dtype=dtype), + ) + monkeypatch.setattr(decoder_runner.torch.cuda, "synchronize", lambda: None) + calls = [] + + def _fake_fused(**kwargs): + calls.append(kwargs["num_real"]) + if kwargs["num_real"] == 32: + raise RuntimeError("synthetic compile failure") + return torch.empty(kwargs["num_real"], dtype=torch.int32) + + monkeypatch.setattr(dspark_schedule, "schedule_verify_lens_topk_fused_fill", _fake_fused) + DecoderRunner._warmup_dspark_fused_scheduler(engine) + assert calls == [16, 32] + assert engine._dspark_fused_schedule_ready_sizes == {16} + + +def test_runtime_failure_retires_only_one_shape_and_retries_once(): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + from tensorrt_llm._torch.speculative.dspark_schedule import DSparkFusedScheduleError + + engine = object.__new__(DecoderRunner) + engine._dspark_fused_schedule_ready_sizes = {16, 32} + calls = [] + + def _select(**kwargs): + calls.append(kwargs["use_fused_exact"]) + if kwargs["use_fused_exact"]: + raise DSparkFusedScheduleError("synthetic launch failure") + return "tensor-result" + + result = DecoderRunner._select_dspark_windows_with_fused_fallback( + engine, + select_fn=_select, + selector_kwargs={}, + padded_bs=16, + ) + assert result == "tensor-result" + assert calls == [True, False] + assert engine._dspark_fused_schedule_ready_sizes == {32} + assert engine._dspark_fused_schedule_failure_counts == {16: 1} + + calls.clear() + result = DecoderRunner._select_dspark_windows_with_fused_fallback( + engine, + select_fn=_select, + selector_kwargs={}, + padded_bs=16, + ) + assert result == "tensor-result" + assert calls == [False] + assert engine._dspark_fused_schedule_failure_counts == {16: 1} + + +def _tensor_prologue_case(): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + + confidence = torch.tensor([[10.0, 10.0, 10.0], [-10.0, -10.0, -10.0]]) + planner = SimpleNamespace( + cfg=DSparkScheduleConfig(block_size=3, min_verify_len=1), + apply_calibration=None, + ) + worker = SimpleNamespace( + verify_planner=planner, + staged_confidence_buffer=lambda: confidence, + batch_slot_view=lambda rows: torch.arange(rows), + verified_draft_seq_cuda=lambda: None, + ) + engine = object.__new__(DecoderRunner) + engine.cuda_graph_runner = SimpleNamespace(agreed_ragged_bucket=6, ragged_pad_verify_len=0) + engine._config = SimpleNamespace(use_mrope=False) + engine._dspark_device_budget = 2 + engine._dspark_prev_covers_batch = True + engine._dspark_last_padded_bs = 2 + engine._dspark_last_num_real = 2 + engine._dspark_fused_schedule_ready_sizes = set() + engine._get_spec_worker = lambda: worker + engine.ragged_verify_lens_cuda = torch.tensor([3, 3], dtype=torch.int32) + engine.ragged_qo_indptr_cuda = torch.tensor([0, 3, 6], dtype=torch.int32) + engine.previous_batch_indices_cuda = torch.tensor([0, 1], dtype=torch.int64) + engine.input_ids_cuda = torch.full((6,), -1, dtype=torch.int32) + engine.position_ids_cuda = torch.tensor([10, 11, 12, 20, 21, 22], dtype=torch.int32) + engine.previous_pos_indices_cuda = torch.full((6,), -1, dtype=torch.int64) + engine.previous_pos_id_offsets_cuda = torch.full((6,), -1, dtype=torch.int32) + engine.previous_kv_lens_offsets_cuda = torch.tensor([-1, -2], dtype=torch.int32) + engine.draft_tokens_cuda = torch.full((4,), -1, dtype=torch.int32) + events = [] + attn_metadata = SimpleNamespace( + num_contexts=0, + kv_lens_cuda=torch.tensor([16, 26], dtype=torch.int32), + apply_device_ragged_layout=lambda lens, owners, correction: events.append( + ("attention", lens.clone(), owners.clone(), correction.clone()) + ), + ) + spec_metadata = SimpleNamespace( + remap_expanded_sampling_params=lambda owners, count: events.append( + ("sampling", owners.clone(), count) + ), + ) + new_tensors = SimpleNamespace( + new_tokens=torch.tensor([[100, 200], [101, 201], [102, 202], [103, 203]]), + new_tokens_lens=torch.tensor([2, 1], dtype=torch.int32), + next_draft_tokens=torch.tensor([[11, 12, 13], [21, 22, 23]]), + ) + return ( + engine, + {"spec_metadata": spec_metadata, "attn_metadata": attn_metadata}, + new_tensors, + events, + ) + + +def test_tensor_prologue_rebuilds_owners_positions_and_both_kv_terms(): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + + engine, inputs, new_tensors, events = _tensor_prologue_case() + lens_address = engine.ragged_verify_lens_cuda.data_ptr() + qo_address = engine.ragged_qo_indptr_cuda.data_ptr() + assert DecoderRunner._apply_device_window_prologue(engine, inputs, new_tensors) + assert engine._dspark_device_budget is None + assert engine.ragged_verify_lens_cuda.tolist() == [4, 2] + assert engine.ragged_qo_indptr_cuda.tolist() == [0, 4, 6] + assert engine.ragged_verify_lens_cuda.data_ptr() == lens_address + assert engine.ragged_qo_indptr_cuda.data_ptr() == qo_address + assert engine.input_ids_cuda.tolist() == [100, 101, 102, 103, 200, 201] + assert engine.position_ids_cuda.tolist() == [10, 11, 12, 13, 20, 21] + assert engine.previous_pos_indices_cuda.tolist() == [0, 0, 0, 0, 1, 1] + assert engine.previous_pos_id_offsets_cuda.tolist() == [2, 2, 2, 2, 1, 1] + assert engine.draft_tokens_cuda.tolist() == [11, 12, 13, 21] + assert inputs["attn_metadata"].kv_lens_cuda.tolist() == [18, 24] + assert engine.previous_kv_lens_offsets_cuda.tolist() == [-2, -1] + assert ( + inputs["attn_metadata"].kv_lens_cuda + engine.previous_kv_lens_offsets_cuda + ).tolist() == [16, 23] + assert [event[0] for event in events] == ["sampling", "attention"] + assert events[0][1].tolist() == [0, 0, 0, 0, 1, 1] + assert events[0][2] == 6 + + +def test_decode_only_prologue_declines_local_prefill_without_mutating_buffers(): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + + engine, inputs, new_tensors, events = _tensor_prologue_case() + inputs["attn_metadata"].num_contexts = 1 + tensors = ( + engine.input_ids_cuda, + engine.position_ids_cuda, + engine.ragged_verify_lens_cuda, + engine.ragged_qo_indptr_cuda, + inputs["attn_metadata"].kv_lens_cuda, + engine.previous_kv_lens_offsets_cuda, + ) + originals = [tensor.clone() for tensor in tensors] + assert not DecoderRunner._apply_device_window_prologue(engine, inputs, new_tensors) + for tensor, original in zip(tensors, originals): + torch.testing.assert_close(tensor, original, rtol=0, atol=0) + assert not events + + +@pytest.mark.parametrize("confidence,fused", [(False, False), (True, False), (True, True)]) +def test_workspace_allocation_follows_decoder_state_initialization(monkeypatch, confidence, fused): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import runner as decoder_runner + + original_empty, original_zeros = torch.empty, torch.zeros + + def cpu_empty(*args, **kwargs): + kwargs["device"] = "cpu" + return original_empty(*args, **kwargs) + + def cpu_zeros(*args, **kwargs): + kwargs["device"] = "cpu" + return original_zeros(*args, **kwargs) + + monkeypatch.setattr(torch, "empty", cpu_empty) + monkeypatch.setattr(torch, "zeros", cpu_zeros) + monkeypatch.setattr(decoder_runner, "prefer_pinned", lambda: False) + monkeypatch.setattr(decoder_runner, "LoraParamBuilder", lambda **kwargs: SimpleNamespace()) + engine = object.__new__(DecoderRunner) + engine._config = SimpleNamespace( + is_spec_decode=True, + max_draft_loop_tokens=3, + max_batch_size=4, + max_num_tokens=16, + max_beam_width=1, + max_seq_len=16, + use_mrope=False, + attention_backend="TRTLLM", + spec_config=SimpleNamespace( + enable_confidence_scheduling=confidence, + enable_fused_confidence_scheduler=fused, + ), + ) + engine._compute_dynamic_draft_len_mapping = lambda: {} + # Follow the constructor order: buffers exist before mutable state/flags. + engine._allocate_decoder_buffers() + engine._init_decoder_state() + if confidence and fused: + workspace = engine._dspark_device_window_workspace + assert workspace.verify_lens.numel() == 5 + assert workspace.req_idx.numel() == 16 + assert workspace.verify_lens.device == engine.input_ids_cuda.device + else: + assert engine._dspark_device_window_workspace is None diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts.py index 04ce5bc4a276..0202110e8d7f 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts.py @@ -11,6 +11,8 @@ resolve_confidence_head, ) +pytestmark = pytest.mark.cpu_only + @pytest.mark.parametrize("key", ["sts_temperatures", "temperatures"]) def test_loads_native_and_sglang_temperature_keys(tmp_path, key): @@ -51,6 +53,14 @@ def test_rejects_non_list_temperatures(tmp_path, key, temperatures): load_sts_temperatures_from_path(str(path)) +@pytest.mark.parametrize("temperature", [0.0, -1.0, float("nan"), float("inf"), -float("inf")]) +def test_rejects_nonfinite_or_nonpositive_temperatures(tmp_path, temperature): + path = tmp_path / "sts.json" + path.write_text(json.dumps({"temperatures": [1.0, temperature]}), encoding="utf-8") + with pytest.raises(ValueError, match="finite positive"): + load_sts_temperatures_from_path(str(path)) + + class _Head: pass diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index 40786627b200..fdfcca0155b2 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -137,7 +137,11 @@ def test_worker_lazy_init_window_buffers(): assert worker._position_initialized.shape == (9,) assert worker._scratch_slot == 8 # Dummy-id floor separates real request ids from CUDA-graph padding ids. - assert worker._graph_dummy_id_floor == (1 << 64) - 1 - worker.max_draft_len + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import cuda_graph_dummy_request_id + + assert worker._graph_dummy_id_floor == cuda_graph_dummy_request_id( + worker.max_draft_len, variant=1, max_draft_len=worker.max_draft_len + ) # The scratch row is never handed out through the free pool. assert list(worker._free_slots) == list(range(8)) assert worker._batch_to_slot is not None @@ -422,17 +426,22 @@ def test_prepare_assigns_slots_to_disagg_generation_requests(): assert worker._batch_to_slot[:2].tolist() == [s0, s1] -def test_prepare_keeps_dummy_generation_requests_on_scratch_row(): +@pytest.mark.parametrize("physical_k", [3, 4, 5]) +@pytest.mark.parametrize("variant", [0, 1]) +def test_prepare_keeps_dummy_generation_requests_on_scratch_row(physical_k, variant): """ADP-idle (id 0) and CUDA-graph padding dummies never consume a real slot.""" - from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import cuda_graph_dummy_request_id from tensorrt_llm._torch.pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID worker = _make_worker() + worker.spec_config.max_draft_len = physical_k meta = _make_metadata(max_num_requests=4) - worker._lazy_init(_fake_draft_model(), meta) + draft_model = _fake_draft_model() + draft_model.block_size = physical_k + worker._lazy_init(draft_model, meta) meta._dspark_worker = worker - graph_dummy = CUDA_GRAPH_DUMMY_REQUEST_ID - worker.max_draft_len + graph_dummy = cuda_graph_dummy_request_id(physical_k, variant=variant, max_draft_len=physical_k) meta.request_ids = [1000, ATTENTION_DP_DUMMY_REQUEST_ID, graph_dummy] meta.num_generations = 3 meta.prepare() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_stamp_wrap_cpu.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_stamp_wrap_cpu.py new file mode 100644 index 000000000000..d7101c14894f --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_stamp_wrap_cpu.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Draft-stamp scalar bounds and neutral reset without device allocation.""" + +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark import DSv4DSparkWorker +from tensorrt_llm._torch.speculative.dspark_schedule import NEUTRAL_CONFIDENCE_LOGIT + +pytestmark = pytest.mark.cpu_only + + +@pytest.mark.parametrize("previous", [0, (1 << 31) - 2, (1 << 31) - 1]) +def test_stamp_sequence_stays_representable_and_wraps_neutrally(previous): + worker = object.__new__(DSv4DSparkWorker) + worker._draft_seq_host = previous + worker._draft_seq_cuda = torch.tensor(previous, dtype=torch.int32) + worker._confidence_stamp = torch.tensor([previous, 0], dtype=torch.int32) + worker._confidence_logits = torch.full((2, 3), -3.0) + pointers = [ + worker._draft_seq_cuda.data_ptr(), + worker._confidence_stamp.data_ptr(), + worker._confidence_logits.data_ptr(), + ] + + assert worker.bump_draft_seq() == previous + current = (previous + 1) & ((1 << 31) - 1) + assert worker._draft_seq_host == worker._draft_seq_cuda.item() == current + assert worker.verified_draft_seq_cuda().item() == current - 1 + if current == 0: + assert worker._confidence_stamp.tolist() == [0, 0] + assert torch.all(worker._confidence_logits == NEUTRAL_CONFIDENCE_LOGIT) + # A real draft emitted at sequence0 is fresh on the next step. Rows + # without that draft remain neutral even if their reset stamp matches. + worker._confidence_logits[0].fill_(-2.0) + assert worker.bump_draft_seq() == 0 + assert worker.verified_draft_seq_cuda().item() == 0 + assert torch.all(worker._confidence_logits[0] == -2.0) + assert torch.all(worker._confidence_logits[1] == NEUTRAL_CONFIDENCE_LOGIT) + else: + assert worker._confidence_stamp.tolist() == [previous, 0] + assert torch.all(worker._confidence_logits == -3.0) + assert pointers == [ + worker._draft_seq_cuda.data_ptr(), + worker._confidence_stamp.data_ptr(), + worker._confidence_logits.data_ptr(), + ] + + +def test_stamp_wrap_without_allocated_confidence_buffers(): + worker = object.__new__(DSv4DSparkWorker) + worker._draft_seq_host = (1 << 31) - 1 + worker._draft_seq_cuda = None + worker._confidence_stamp = None + worker._confidence_logits = None + assert worker.bump_draft_seq() == (1 << 31) - 1 + assert worker._draft_seq_host == 0 + assert worker.verified_draft_seq_cuda() is None diff --git a/tests/unittest/_torch/speculative/test_device_window_sampling_layout.py b/tests/unittest/_torch/speculative/test_device_window_sampling_layout.py new file mode 100644 index 000000000000..bcab9063a968 --- /dev/null +++ b/tests/unittest/_torch/speculative/test_device_window_sampling_layout.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import types + +import torch + +from tensorrt_llm._torch.speculative.interface import _RNG_SLOT_SPAN, SpecMetadata + + +def test_device_window_repacking_follows_fresh_token_owners(): + owners = torch.tensor([0, 0, 1, 2, 2, 2], dtype=torch.long) + meta = types.SimpleNamespace( + request_temperatures=torch.tensor([0.2, 0.7, 1.1]), + temperatures=torch.empty(6), + request_top_ks=torch.tensor([2, 7, 11], dtype=torch.int32), + top_ks=torch.empty(6, dtype=torch.int32), + request_top_ps=torch.tensor([0.3, 0.8, 0.95]), + top_ps=torch.empty(6), + request_min_ps=torch.tensor([0.1, 0.2, 0.4]), + min_ps=torch.empty(6), + request_seeds=torch.tensor([101, 202, 303], dtype=torch.int64), + seeds=torch.empty(6, dtype=torch.int64), + request_offsets=torch.tensor([10, 20, 30], dtype=torch.int64), + offsets=torch.empty(6, dtype=torch.int64), + qo_indptr=torch.tensor([0, 2, 3, 6], dtype=torch.int32), + _sampling_params_signature=[("request",), ("expanded",)], + ) + + SpecMetadata.remap_expanded_sampling_params(meta, owners, owners.numel()) + + assert torch.equal(meta.temperatures, meta.request_temperatures.index_select(0, owners)) + assert torch.equal(meta.top_ks, meta.request_top_ks.index_select(0, owners)) + assert torch.equal(meta.top_ps, meta.request_top_ps.index_select(0, owners)) + assert torch.equal(meta.min_ps, meta.request_min_ps.index_select(0, owners)) + assert torch.equal(meta.seeds, meta.request_seeds.index_select(0, owners)) + row_slots = torch.tensor([0, 1, 0, 0, 1, 2], dtype=torch.int64) + expected_offsets = meta.request_offsets.index_select(0, owners) + row_slots * _RNG_SLOT_SPAN + assert torch.equal(meta.offsets, expected_offsets) + assert meta._sampling_params_signature[0] == ("request",) + assert meta._sampling_params_signature[1] is None + + +def test_device_window_repacking_ignores_unallocated_optional_buffers(): + meta = types.SimpleNamespace( + request_temperatures=None, + temperatures=None, + request_top_ks=None, + top_ks=None, + request_top_ps=None, + top_ps=None, + request_min_ps=None, + min_ps=None, + request_seeds=None, + seeds=None, + request_offsets=None, + offsets=None, + _sampling_params_signature=[None, None], + ) + SpecMetadata.remap_expanded_sampling_params(meta, torch.tensor([0]), 1) + + +def test_device_window_repacking_forces_next_host_expansion(): + values = ((0.2, 2, 0.3, 0.1), (0.7, 7, 0.8, 0.2)) + per_request = [(0.2, 2, 0.3, 0.1, 2), (0.7, 7, 0.8, 0.2, 4)] + meta = types.SimpleNamespace( + request_temperatures=torch.tensor([0.2, 0.7]), + temperatures=torch.empty(6), + request_top_ks=torch.tensor([2, 7], dtype=torch.int32), + top_ks=torch.empty(6, dtype=torch.int32), + request_top_ps=torch.tensor([0.3, 0.8]), + top_ps=torch.empty(6), + request_min_ps=torch.tensor([0.1, 0.2]), + min_ps=torch.empty(6), + request_seeds=None, + seeds=None, + request_offsets=None, + offsets=None, + _sampling_params_signature=[values, (values, (2, 4))], + ) + SpecMetadata.remap_expanded_sampling_params(meta, torch.tensor([0, 1, 1, 1, 0, 0]), 6) + + need_request, need_expanded = SpecMetadata._sampling_params_buffers_need_update( + meta, per_request + ) + assert need_request is False + assert need_expanded is True diff --git a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py index 0c9468c48a75..6bb25b0734e3 100644 --- a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py +++ b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py @@ -16,6 +16,7 @@ from tensorrt_llm._utils import get_sm_version _ROPE_DIM = 64 +_PHYSICAL_KS = (1, 2, 3, 4, 5, 6, 7, 8) pytestmark = pytest.mark.skipif( not torch.cuda.is_available() @@ -106,7 +107,7 @@ def _reference( ("invalid_case", "reason"), [ ("q_dtype", "q dtype must be BF16"), - ("block_size", "draft block size must be 5 or 6"), + ("block_size", "draft block size must be 1 through 8"), ("valid_len_dtype", "valid_len must use INT64"), ("freqs_shape", "inverse_rope_freqs shape must be"), ], @@ -119,8 +120,8 @@ def test_fused_dsv4_dspark_attention_support_gate_logs_rejection(monkeypatch, in if invalid_case == "q_dtype": inputs[0] = q.float() elif invalid_case == "block_size": - inputs[0] = q[:, :4].contiguous() - inputs[4] = freqs[:, :4].contiguous() + inputs[0] = q.new_empty((q.shape[0], 9, *q.shape[2:])) + inputs[4] = freqs.new_empty((q.shape[0], 9, *freqs.shape[2:])) elif invalid_case == "valid_len_dtype": inputs[2] = inputs[2].to(torch.int32) else: @@ -177,7 +178,7 @@ def test_fused_dsv4_dspark_attention_rejects_invalid_inputs_before_launch( ) -@pytest.mark.parametrize("block", (5, 6)) +@pytest.mark.parametrize("block", _PHYSICAL_KS) @pytest.mark.parametrize( ("start_pos_values", "valid_len_values"), [ @@ -188,8 +189,10 @@ def test_fused_dsv4_dspark_attention_rejects_invalid_inputs_before_launch( ([5, 100], [6, 101]), # Bootstrapped positions with short physical suffixes exercise wraparound. ([257, 390], [3, 5]), + # The first row is a graph-padding request with no initialized context. + ([0, 5], [0, 6]), ], - ids=["full_window", "partial_window", "wrapped_suffix"], + ids=["full_window", "partial_window", "wrapped_suffix", "empty_scratch_window"], ) def test_fused_dsv4_dspark_attention_matches_reference(block, start_pos_values, valid_len_values): from tensorrt_llm._torch.custom_ops.dspark_attention_custom_op import ( @@ -199,6 +202,9 @@ def test_fused_dsv4_dspark_attention_matches_reference(block, start_pos_values, q, main_kv, block_kv, kv_cache, slots, start_pos, sink, blk_freqs, inverse_rope_freqs = ( _make_inputs(29, block=block, start_pos_values=start_pos_values) ) + if valid_len_values[0] == 0: + # Use a reserved scratch page rather than aliasing a live request slot. + slots[0] = kv_cache.shape[0] - 1 valid_len = torch.tensor(valid_len_values, device=q.device, dtype=torch.long) expected, expected_cache = _reference( q, @@ -215,6 +221,8 @@ def test_fused_dsv4_dspark_attention_matches_reference(block, start_pos_values, draft_block, slots_i32, cache_seqs = _prepare_attention_inputs( main_kv, block_kv, kv_cache, slots, start_pos ) + # A wrong draft mask must not accidentally pass because padding is zero. + draft_block[:, block:].fill_(17) actual = fused_dsv4_dspark_attention( q, @@ -232,7 +240,8 @@ def test_fused_dsv4_dspark_attention_matches_reference(block, start_pos_values, torch.testing.assert_close(kv_cache, expected_cache, rtol=0, atol=0) -def test_fused_dsv4_dspark_attention_cuda_graph_replay(): +@pytest.mark.parametrize("block", _PHYSICAL_KS) +def test_fused_dsv4_dspark_attention_cuda_graph_replay(block): from tensorrt_llm._torch.custom_ops.dspark_attention_custom_op import ( fused_dsv4_dspark_attention, warmup_fused_dsv4_dspark_attention, @@ -244,9 +253,9 @@ def test_fused_dsv4_dspark_attention_cuda_graph_replay(): ) q, main_kv, block_kv, kv_cache, slots, start_pos, sink, blk_freqs, inverse_rope_freqs = ( - _make_inputs(3, start_pos_values=[257, 5, 390]) + _make_inputs(3, block=block, start_pos_values=[257, 5, 390, 0]) ) - batch, block, dim = block_kv.shape + batch, _, dim = block_kv.shape scale = dim**-0.5 main_x = main_kv.unsqueeze(1) block_x = block_kv @@ -256,8 +265,10 @@ def test_fused_dsv4_dspark_attention_cuda_graph_replay(): main_freqs[..., 0] = 1 block_freqs[..., 0] = 1 slots = slots.long() + slots[-1] = kv_cache.shape[0] - 1 start_pos = start_pos.long() valid_len = _legacy_valid_len(start_pos) + valid_len[-1] = 0 # The named best-effort prewarm calls the same self-JIT ops used below, # so graph capture sees only hot compile-cache entries. @@ -282,7 +293,15 @@ def test_fused_dsv4_dspark_attention_cuda_graph_replay(): main_x.copy_(torch.randn_like(main_x)) block_x.copy_(torch.randn_like(block_x)) - valid_len.copy_(torch.tensor([3, 2, 7], device=q.device)) + q.copy_(torch.randn_like(q)) + sink.copy_(torch.randn_like(sink)) + # Replay must read new positions, page mappings, and inverse-RoPE phases. + start_pos.copy_(torch.tensor([390, 127, 257, 0], device=q.device)) + slots[:3].copy_(slots[:3].roll(1)) + valid_len.copy_(torch.tensor([3, 128, 7, 0], device=q.device)) + table = precompute_dspark_freqs_cis(_ROPE_DIM, 512, device=q.device) + blk_freqs = table[start_pos.unsqueeze(1) + 1 + torch.arange(block, device=q.device)] + inverse_rope_freqs.copy_(torch.view_as_real(blk_freqs)) expected_main = cute_dsl_dspark_rmsnorm_rope( main_x, weight, main_freqs, 1, _ROPE_DIM, 1e-6, True, True, False ).squeeze(1) @@ -312,13 +331,14 @@ def test_fused_dsv4_dspark_attention_cuda_graph_replay(): @pytest.mark.parametrize("persist", (False, True)) -def test_dspark_attention_forward_batched_matches_fallback(monkeypatch, persist): +@pytest.mark.parametrize("block", _PHYSICAL_KS) +def test_dspark_attention_forward_batched_matches_fallback(monkeypatch, persist, block): import tensorrt_llm._torch.models.modeling_dspark as dspark_attention torch.manual_seed(17) device = torch.device("cuda") dtype = torch.bfloat16 - batch, block, hidden = 2, 6, 64 + batch, hidden = 2, 64 heads, head_dim, rope_dim = 128, 512, 64 q_rank, groups, o_rank, window = 1024, 8, 32, 128 @@ -444,7 +464,8 @@ def test_warmup_is_best_effort(monkeypatch): assert "synthetic compile failure" in warnings[0] -def test_self_jit_reuses_one_kernel_across_runtime_shapes_and_scales(): +@pytest.mark.parametrize("block", _PHYSICAL_KS) +def test_self_jit_reuses_one_kernel_across_runtime_shapes_and_scales(block): from tensorrt_llm._torch.custom_ops.dspark_attention_custom_op import ( _dspark_attention_kernel_cache, _get_dspark_arch_str, @@ -459,13 +480,13 @@ def test_self_jit_reuses_one_kernel_across_runtime_shapes_and_scales(): assert _get_dspark_arch_str(101) is None assert _get_dspark_arch_str(109) is None assert _get_dspark_arch_str() in ("sm_100", "sm_103") - assert is_dsv4_dspark_attention_config_supported(5, 128, 512, 128) - assert not is_dsv4_dspark_attention_config_supported(4, 128, 512, 128) + assert is_dsv4_dspark_attention_config_supported(block, 128, 512, 128) + assert not is_dsv4_dspark_attention_config_supported(9, 128, 512, 128) assert not is_dsv4_dspark_attention_config_supported(5, 24, 512, 128) _dspark_attention_kernel_cache.clear() q, main_kv, block_kv, kv_cache, slots, start_pos, sink, _, freqs = _make_inputs( - 11, start_pos_values=[300, 4, 250], cache_pages=40 + 11, block=block, start_pos_values=[300, 4, 250], cache_pages=40 ) scale = q.shape[-1] ** -0.5 valid_len = _legacy_valid_len(start_pos) @@ -505,7 +526,9 @@ def test_self_jit_reuses_one_kernel_across_runtime_shapes_and_scales(): # hit on the same compiled object without compiler calls or runtime padding. for batch in (1, 3, 32): values = [5 + (37 * i) % 386 for i in range(batch)] - args = _make_inputs(17 + batch, start_pos_values=values, cache_pages=batch + 41) + args = _make_inputs( + 17 + batch, block=block, start_pos_values=values, cache_pages=batch + 41 + ) q_b, main_b, block_b, cache_b, slots_b, pos_b, sink_b, blk_freqs_b, freqs_b = args if batch == 3: cache_b = cache_b.clone() @@ -543,6 +566,43 @@ def test_self_jit_reuses_one_kernel_across_runtime_shapes_and_scales(): assert len(_dspark_attention_kernel_cache) == 1 +def test_attention_cache_isolates_all_k_with_shared_dummy_scratch(): + import tensorrt_llm._torch.custom_ops.dspark_attention_custom_op as attention_op + + attention_op._dspark_attention_kernel_cache.clear() + for block in _PHYSICAL_KS: + attention_op.warmup_fused_dsv4_dspark_attention(block, 1e-6) + arch = attention_op._get_dspark_arch_str() + kernels = attention_op._dspark_attention_kernel_cache.copy() + assert set(kernels) == {(block, arch) for block in _PHYSICAL_KS} + assert len({id(kernel) for kernel in kernels.values()}) == len(_PHYSICAL_KS) + + # Keep all specializations live while alternating K. Two padding requests + # read the reserved scratch page, never a real request's window. No duplicate + # cache-write launches are needed for this attention-only aliasing check. + for block in reversed(_PHYSICAL_KS): + q, main, block_kv, cache, slots, pos, sink, blk_freqs, freqs = _make_inputs( + 59, block=block, start_pos_values=[257, 0, 0] + ) + scratch = cache.shape[0] - 1 + slots.copy_(torch.tensor([0, scratch, scratch], device=q.device, dtype=slots.dtype)) + main[1:].zero_() + cache[0, 257 % 128].copy_(main[0]) + cache[scratch, 0].zero_() + valid = torch.tensor([3, 0, 0], device=q.device, dtype=torch.long) + draft = block_kv.new_zeros((3, 8, 512)) + draft[:, :block].copy_(block_kv) + expected, expected_cache = _reference( + q, main, block_kv, cache, slots, pos, valid, sink, blk_freqs, 512**-0.5 + ) + actual = attention_op.fused_dsv4_dspark_attention( + q, draft, cache, slots, pos, valid, sink, freqs, 512**-0.5 + ) + torch.testing.assert_close(actual, expected, rtol=5e-2, atol=3e-2) + torch.testing.assert_close(cache, expected_cache, rtol=0, atol=0) + assert attention_op._dspark_attention_kernel_cache == kernels + + def test_compile_without_real_specimens_and_cached_wrapper_avoids_views(monkeypatch): import tensorrt_llm._torch.custom_ops.dspark_attention_custom_op as dspark_attention_op @@ -596,11 +656,12 @@ def unexpected_host_op(*args, **kwargs): assert output.shape == q.shape -def test_attention_cache_miss_rejects_cuda_graph_capture(monkeypatch): +@pytest.mark.parametrize("block", _PHYSICAL_KS) +def test_attention_cache_miss_rejects_cuda_graph_capture(monkeypatch, block): import tensorrt_llm._torch.custom_ops.dspark_attention_custom_op as dspark_attention_op dspark_attention_op._dspark_attention_kernel_cache.clear() - q, main_kv, block_kv, kv_cache, slots, start_pos, sink, _, freqs = _make_inputs(43, block=5) + q, main_kv, block_kv, kv_cache, slots, start_pos, sink, _, freqs = _make_inputs(43, block=block) draft_block, slots_i32, cache_seqs = _prepare_attention_inputs( main_kv, block_kv, kv_cache, slots, start_pos ) @@ -630,7 +691,8 @@ def test_attention_cache_miss_rejects_cuda_graph_capture(monkeypatch): ) -def test_preparation_cache_misses_reject_cuda_graph_capture(monkeypatch): +@pytest.mark.parametrize("block", _PHYSICAL_KS) +def test_preparation_cache_misses_reject_cuda_graph_capture(monkeypatch, block): import tensorrt_llm._torch.custom_ops.dspark_rmsnorm_rope_custom_op as preparation_op preparation_op._compile_dspark_rmsnorm_rope_cache_write.cache_clear() @@ -645,10 +707,11 @@ def test_preparation_cache_misses_reject_cuda_graph_capture(monkeypatch): with pytest.raises(RuntimeError, match="cache-write must be warmed up"): preparation_op._compile_dspark_rmsnorm_rope_cache_write(1e-6) with pytest.raises(RuntimeError, match="draft-block must be warmed up"): - preparation_op._compile_dspark_rmsnorm_rope_draft_block(5, 1e-6) + preparation_op._compile_dspark_rmsnorm_rope_draft_block(block, 1e-6) -def test_preparation_self_jit_covers_dynamic_batches(): +@pytest.mark.parametrize("block", _PHYSICAL_KS) +def test_preparation_self_jit_covers_dynamic_batches(block): from tensorrt_llm._torch.custom_ops.dspark_rmsnorm_rope_custom_op import ( _compile_dspark_rmsnorm_rope_cache_write, _compile_dspark_rmsnorm_rope_draft_block, @@ -659,7 +722,7 @@ def test_preparation_self_jit_covers_dynamic_batches(): _compile_dspark_rmsnorm_rope_cache_write.cache_clear() _compile_dspark_rmsnorm_rope_draft_block.cache_clear() q, main_kv, block_kv, kv_cache, slots, start_pos, _, _, _ = _make_inputs( - 23, start_pos_values=[300, 4, 250], cache_pages=40 + 23, block=block, start_pos_values=[300, 4, 250], cache_pages=40 ) weight = torch.ones(512, device=q.device, dtype=q.dtype) main_freqs = torch.zeros(q.shape[0], _ROPE_DIM // 2, 2, device=q.device) @@ -668,7 +731,7 @@ def test_preparation_self_jit_covers_dynamic_batches(): block_freqs[..., 0] = 1 for batch in (1, 3, 32): - args = _make_inputs(31 + batch, batch=batch, cache_pages=40) + args = _make_inputs(31 + batch, block=block, batch=batch, cache_pages=40) q_b, main_b, block_b, cache_b, slots_b, pos_b, _, _, _ = args main_freqs_b = torch.zeros(batch, _ROPE_DIM // 2, 2, device=q.device) block_freqs_b = torch.zeros(batch * q_b.shape[1], _ROPE_DIM // 2, 2, device=q.device) diff --git a/tests/unittest/llmapi/test_dspark_workload_admission_config.py b/tests/unittest/llmapi/test_dspark_workload_admission_config.py new file mode 100644 index 000000000000..bf144525a4d0 --- /dev/null +++ b/tests/unittest/llmapi/test_dspark_workload_admission_config.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Startup-boundary tests for DSpark workload admission.""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from tensorrt_llm.llmapi.llm_args import CudaGraphConfig, DSparkDecodingConfig, TorchLlmArgs + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("outcome", ["missing", "negative", "stale"]) +def test_dspark_workload_admission_disables_before_runtime_construction(tmp_path, outcome): + (tmp_path / "config.json").write_text('{"dspark_block_size": 2}') + spec_cfg = DSparkDecodingConfig( + max_draft_len=2, + speculative_model=str(tmp_path), + enable_confidence_scheduling=True, + validate_confidence_artifact_freshness=True, + confidence_sps_table_path=(None if outcome == "missing" else "/sealed/sps.json"), + confidence_sps_live_fingerprint_path=( + None if outcome == "missing" else "/sealed/live.json" + ), + confidence_admission_receipt_path=( + None if outcome == "missing" else "/sealed/admission.json" + ), + confidence_admission_receipt_sha256=(None if outcome == "missing" else "a" * 64), + ) + if outcome == "negative": + patch_kwargs = {"return_value": SimpleNamespace(admitted=False)} + elif outcome == "stale": + patch_kwargs = {"side_effect": ValueError("stale receipt")} + else: + patch_kwargs = {"side_effect": TypeError("missing receipt")} + with patch( + "tensorrt_llm._torch.speculative.dspark_planner.load_confidence_workload_admission", + **patch_kwargs, + ): + # No CUDA-graph configuration is supplied. Construction succeeds only + # if admission resolves the feature off before the confidence-specific + # environment validator observes it. + args = TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + ) + + assert args.speculative_config.max_draft_len == 2 + assert args.speculative_config.enable_confidence_scheduling is False + assert args.speculative_config.confidence_sps_table_path is None + assert args.speculative_config.confidence_sps_live_fingerprint_path is None + assert args.speculative_config.confidence_admission_receipt_path is None + assert args.speculative_config.confidence_admission_receipt_sha256 is None + + +@pytest.mark.cpu_only +def test_dspark_positive_workload_admission_preserves_confidence_path(tmp_path): + (tmp_path / "config.json").write_text('{"dspark_block_size": 2}') + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"mtp.0.weight": "model.safetensors"}}' + ) + spec_cfg = DSparkDecodingConfig( + max_draft_len=2, + speculative_model=str(tmp_path), + enable_confidence_scheduling=True, + validate_confidence_artifact_freshness=True, + confidence_sps_table_path="/sealed/sps.json", + confidence_sps_live_fingerprint_path="/sealed/live.json", + confidence_admission_receipt_path="/sealed/admission.json", + confidence_admission_receipt_sha256="a" * 64, + ) + with patch( + "tensorrt_llm._torch.speculative.dspark_planner.load_confidence_workload_admission", + return_value=SimpleNamespace(admitted=True), + ): + args = TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + cuda_graph_config=CudaGraphConfig(batch_sizes=[1], enable_padding=True), + ) + + assert args.speculative_config.enable_confidence_scheduling is True + assert args.speculative_config.confidence_admission_receipt_path == "/sealed/admission.json" + assert args.speculative_config.confidence_admission_receipt_sha256 == "a" * 64 + + +@pytest.mark.cpu_only +def test_dspark_reuse_keeps_enabled_without_workload_admission(tmp_path): + (tmp_path / "config.json").write_text('{"dspark_block_size": 2}') + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"mtp.0.weight": "model.safetensors"}}' + ) + spec_cfg = DSparkDecodingConfig( + max_draft_len=2, + speculative_model=str(tmp_path), + enable_confidence_scheduling=True, + confidence_sps_table_path="/original/sps.json", + confidence_sts_path="/original/sts.json", + ) + with patch( + "tensorrt_llm._torch.speculative.dspark_planner.load_confidence_workload_admission", + side_effect=AssertionError("Reuse must not construct an admission receipt"), + ): + args = TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + cuda_graph_config=CudaGraphConfig(batch_sizes=[1], enable_padding=True), + ) + + assert args.speculative_config.enable_confidence_scheduling is True + assert args.speculative_config.validate_confidence_artifact_freshness is False + assert args.speculative_config.confidence_admission_receipt_path is None + assert args.speculative_config.confidence_admission_receipt_sha256 is None + assert args.speculative_config.confidence_sps_table_path == "/original/sps.json" + assert args.speculative_config.confidence_sts_path == "/original/sts.json" + + +@pytest.mark.cpu_only +def test_dspark_reuse_preserves_padded_graph_requirement(tmp_path): + (tmp_path / "config.json").write_text('{"dspark_block_size": 2}') + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"mtp.0.weight": "model.safetensors"}}' + ) + spec_cfg = DSparkDecodingConfig( + max_draft_len=2, + speculative_model=str(tmp_path), + enable_confidence_scheduling=True, + confidence_sps_table_path="/original/sps.json", + ) + with pytest.raises(ValueError, match="cuda_graph_config"): + TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + cuda_graph_config=None, + ) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 1c720ea3a246..119c04694c0a 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -384,7 +384,7 @@ def test_llm_args_with_kvcache_config(self): llm_args_dict = self._yaml_to_dict(yaml_content) llm_args = TorchLlmArgs(model=llama_model_path, **llm_args_dict) - assert llm_args.kv_cache_config.enable_block_reuse == True + assert llm_args.kv_cache_config.enable_block_reuse assert llm_args.kv_cache_config.max_tokens == 1024 assert llm_args.kv_cache_config.max_attention_window == [ 1024, 1024, 1024 @@ -1198,7 +1198,7 @@ def test_KvCacheConfig_declaration(): attention_dp_events_gather_period_ms=10) pybind_config = config._to_pybind() - assert pybind_config.enable_block_reuse == True + assert pybind_config.enable_block_reuse assert pybind_config.max_tokens == 1024 assert pybind_config.max_attention_window == [1024, 1024, 1024] assert pybind_config.free_gpu_memory_fraction == 0.5 @@ -1234,8 +1234,8 @@ def test_KvCacheConfig_declaration(): ).kv_cache_event_hash_algo == "v2_sha256" assert KvCacheConfig(kv_cache_event_hash_algo="v2_sha256_64" ).kv_cache_event_hash_algo == "v2_sha256_64" - assert pybind_config.enable_partial_reuse == True - assert pybind_config.copy_on_partial_reuse == True + assert pybind_config.enable_partial_reuse + assert pybind_config.copy_on_partial_reuse assert pybind_config.attention_dp_events_gather_period_ms == 10 assert BlockReuseConfig( policy="per_conversation").policy == "per_conversation" @@ -1864,8 +1864,8 @@ def test_DynamicBatchConfig_declaration(): pybind_config = PybindMirror.maybe_to_pybind(config) - assert pybind_config.enable_batch_size_tuning == True - assert pybind_config.enable_max_num_tokens_tuning == True + assert pybind_config.enable_batch_size_tuning + assert pybind_config.enable_max_num_tokens_tuning assert pybind_config.dynamic_batch_moving_average_window == 10 @@ -3094,7 +3094,7 @@ def test_dynamic_batch_config_arbitrary_args(self): config = DynamicBatchConfig(enable_batch_size_tuning=True, enable_max_num_tokens_tuning=True, dynamic_batch_moving_average_window=8) - assert config.enable_batch_size_tuning == True + assert config.enable_batch_size_tuning # Arbitrary arguments should be rejected with pytest.raises( @@ -3139,7 +3139,7 @@ def test_kv_cache_config_arbitrary_args(self): """Test that KvCacheConfig rejects arbitrary arguments.""" # Valid arguments should work config = KvCacheConfig(enable_block_reuse=True, max_tokens=1024) - assert config.enable_block_reuse == True + assert config.enable_block_reuse assert config.max_tokens == 1024 # Arbitrary arguments should be rejected @@ -3154,8 +3154,8 @@ def test_extended_runtime_perf_knob_config_arbitrary_args(self): # Valid arguments should work config = ExtendedRuntimePerfKnobConfig(multi_block_mode=True, cuda_graph_mode=False) - assert config.multi_block_mode == True - assert config.cuda_graph_mode == False + assert config.multi_block_mode + assert not config.cuda_graph_mode # Arbitrary arguments should be rejected with pytest.raises( @@ -3270,8 +3270,8 @@ def test_torch_compile_config_arbitrary_args(self): # Valid arguments should work config = TorchCompileConfig(enable_fullgraph=True, enable_inductor=False) - assert config.enable_fullgraph == True - assert config.enable_inductor == False + assert config.enable_fullgraph + assert not config.enable_inductor # Arbitrary arguments should be rejected with pytest.raises( @@ -3757,8 +3757,8 @@ def test_model_defaults_application(self): assert self.get_model_defaults_called modified_args = llm._executor.engine.model_engine.llm_args - assert modified_args.enable_chunked_prefill == True - assert modified_args.kv_cache_config.enable_block_reuse == False + assert modified_args.enable_chunked_prefill + assert not modified_args.kv_cache_config.enable_block_reuse assert modified_args.kv_cache_config.free_gpu_memory_fraction == 0.75 @pytest.mark.part0 @@ -3778,10 +3778,10 @@ def test_user_overrides_respected(self): assert self.get_model_defaults_called modified_args = llm._executor.engine.model_engine.llm_args - assert modified_args.enable_chunked_prefill == False + assert not modified_args.enable_chunked_prefill assert modified_args.max_batch_size == 42 assert modified_args.max_input_len == 256 - assert modified_args.kv_cache_config.enable_block_reuse == True + assert modified_args.kv_cache_config.enable_block_reuse @pytest.mark.part0 def test_partial_user_override(self): @@ -3798,8 +3798,8 @@ def test_partial_user_override(self): modified_args = llm._executor.engine.model_engine.llm_args assert modified_args.max_batch_size == 42 - assert modified_args.enable_chunked_prefill == True - assert modified_args.kv_cache_config.enable_block_reuse == False + assert modified_args.enable_chunked_prefill + assert not modified_args.kv_cache_config.enable_block_reuse @pytest.mark.part0 def test_empty_nested_config_preserves_defaults(self): @@ -3825,7 +3825,7 @@ def test_empty_nested_config_preserves_defaults(self): # Model defaults set enable_block_reuse=False and # free_gpu_memory_fraction=0.75. An empty KvCacheConfig() # should not prevent these from being applied. - assert modified_args.kv_cache_config.enable_block_reuse == False + assert not modified_args.kv_cache_config.enable_block_reuse assert modified_args.kv_cache_config.free_gpu_memory_fraction == 0.75 @@ -3896,8 +3896,7 @@ def _get_qualified_name(cls: type) -> str: @pytest.mark.cpu_only class TestPydanticBestPractices: - """Ensure that the user-facing LlmArgs and its subfields follow Pydantic best practices. - """ + """Check user-facing LlmArgs and subfields for Pydantic best practices.""" # Fields exempt from Pydantic compatibility checks due to typing limitations or other edge cases. # Avoid adding to this list unless absolutely necessary, especially if a field is user-facing. @@ -5242,3 +5241,154 @@ def test_deepseek_resolves_auto_to_python_on_nixl(self) -> None: cfg = self._pretrained_config(["DeepseekV3ForCausalLM"], "deepseek_v3") _resolve_transceiver_runtime_auto(args, DeepseekV3ForCausalLM, cfg) assert args.cache_transceiver_config.transceiver_runtime == "PYTHON" + + +class TestDSparkConfidenceScheduling: + """Config surface for DSpark confidence-scheduled verification.""" + + def _cfg(self, **kw): + from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig + kw.setdefault("max_draft_len", 5) + return DSparkDecodingConfig(**kw) + + def _scheduled(self, **kw): + """Scheduling on with the required profiled cost table.""" + kw["enable_confidence_scheduling"] = True + kw.setdefault("confidence_sps_table_path", "/tmp/sps.json") + kw.setdefault("confidence_sps_live_fingerprint_path", + "/tmp/runtime.json") + return self._cfg(**kw) + + def test_defaults_are_off(self): + c = self._cfg() + assert c.enable_confidence_scheduling is False + assert c.enable_fused_confidence_scheduler is False + + def test_knobs_require_the_master_switch(self): + for kw in ({ + "confidence_sts_path": "/tmp/sts.json" + }, { + "confidence_sps_table_path": "/tmp/sps.json" + }, { + "confidence_sps_live_fingerprint_path": "/tmp/runtime.json" + }, { + "confidence_admission_receipt_path": "/tmp/admission.json" + }, { + "confidence_admission_receipt_sha256": "a" * 64 + }): + with pytest.raises(ValueError, + match="enable_confidence_scheduling"): + self._cfg(**kw) + + def test_admission_artifacts_are_resolved_at_torch_args_boundary(self): + c = self._cfg(enable_confidence_scheduling=True) + assert c.enable_confidence_scheduling is True + + def test_removed_legacy_verify_len_tiers_are_rejected(self): + with pytest.raises(ValidationError, + match="confidence_verify_len_tiers"): + self._scheduled(confidence_verify_len_tiers=[1, 3, 5]) + + def test_physical_k1_resolves_confidence_scheduling_to_static(self): + c = self._scheduled( + max_draft_len=1, + confidence_sts_path="/tmp/sts.json", + confidence_admission_receipt_path="/tmp/admission.json", + confidence_admission_receipt_sha256="a" * 64, + enable_fused_confidence_scheduler=True, + ) + + assert c.enable_confidence_scheduling is False + assert c.enable_fused_confidence_scheduler is False + assert c.confidence_sts_path is None + assert c.confidence_sps_table_path is None + assert c.confidence_sps_live_fingerprint_path is None + assert c.confidence_admission_receipt_path is None + assert c.confidence_admission_receipt_sha256 is None + + def test_physical_k1_needs_no_confidence_artifacts(self): + c = self._cfg(max_draft_len=1, enable_confidence_scheduling=True) + + assert c.enable_confidence_scheduling is False + assert c.max_total_draft_tokens == 1 + assert c.tokens_per_gen_step == 2 + + def test_physical_k1_remains_valid_when_scheduling_is_off(self): + c = self._cfg(max_draft_len=1) + assert c.enable_confidence_scheduling is False + assert c.max_draft_len == 1 + assert c.max_total_draft_tokens == 1 + assert c.tokens_per_gen_step == 2 + + def test_fused_scheduler_is_independently_gated(self): + with pytest.raises(ValueError, + match="enable_fused_confidence_scheduler"): + self._cfg(enable_fused_confidence_scheduler=True) + c = self._scheduled(enable_fused_confidence_scheduler=True) + assert c.enable_fused_confidence_scheduler is True + + @staticmethod + def _environment(**overrides): + values = dict( + cuda_graph_config=SimpleNamespace(enable_padding=True), + speculative_config=SimpleNamespace( + draft_is_embedded_in_target=True), + guided_decoding_backend=None, + enable_lora=False, + lora_config=None, + sparse_attention_config=None, + pipeline_parallel_size=1, + ) + values.update(overrides) + return SimpleNamespace(**values) + + def test_confidence_environment_rejects_standalone_dspark(self): + environment = self._environment(speculative_config=SimpleNamespace( + draft_is_embedded_in_target=False)) + with pytest.raises(ValueError, + match="only the embedded DeepSeek-V4 draft"): + TorchLlmArgs._validate_dspark_confidence_environment(environment) + + @pytest.mark.parametrize( + ("overrides", "message"), + [ + ({ + "cuda_graph_config": None + }, "requires cuda_graph_config"), + ( + { + "cuda_graph_config": SimpleNamespace(enable_padding=False) + }, + "enable_padding=True", + ), + ({ + "guided_decoding_backend": "xgrammar" + }, "guided decoding"), + ({ + "enable_lora": True + }, "does not yet support LoRA"), + ({ + "pipeline_parallel_size": 2 + }, "does not yet support pipeline parallelism"), + ( + { + "sparse_attention_config": + SimpleNamespace( + enable_heuristic_topk=True, + use_cute_dsl_topk=False, + use_cute_dsl_paged_mqa_logits=False, + ) + }, + "enable_heuristic_topk", + ), + ], + ) + def test_confidence_environment_rejects_unsupported_runtime_combinations( + self, overrides, message): + environment = self._environment(**overrides) + with pytest.raises(ValueError, match=message): + TorchLlmArgs._validate_dspark_confidence_environment(environment) + + def test_confidence_environment_accepts_padded_cuda_graphs(self): + TorchLlmArgs._validate_dspark_confidence_environment( + self._environment())