diff --git a/examples/models/core/deepseek_v4/CONFIDENCE_REPRODUCTION.md b/examples/models/core/deepseek_v4/CONFIDENCE_REPRODUCTION.md new file mode 100644 index 000000000000..192a597a1d91 --- /dev/null +++ b/examples/models/core/deepseek_v4/CONFIDENCE_REPRODUCTION.md @@ -0,0 +1,225 @@ + + + +# DSpark confidence scheduling and reproduction + +Use this guide to run with supplied artifacts or generate your own. Ordinary +inference and benchmarking use the normal TensorRT LLM runtime YAML, APIs and +CLI; the optional reproduction helper is not a serving prerequisite. + +Confidence scheduling is experimental and **off by default**. Explicitly enabling +it retains full physical draft proposals and normal target verification and +sampling, while allowing a smaller measured verification layout. + +## 1. Prepare the runtime and inputs + +Follow the [DeepSeek-V4 Quick Start](README.md#quick-start) and the ordinary +[installation/build instructions](../../../../docs/source/installation/build-from-source.md). +Use a compatible runtime and a supported embedded DSpark checkpoint containing +trained confidence-head weights. Temperature calibration does not train those +weights. No particular cluster is required. + +Choose the following inputs for your workload: + +| Input | Meaning | +| --- | --- | +| Physical `K` | YAML `speculative_config.max_draft_len`, and offline/helper `--physical-k`. It must match the loaded checkpoint/runtime. Compact scheduling needs K > 1; K1 resolves to feature-off static execution. | +| Graph size `G` | Rank-local request count in `cuda_graph_config.batch_sizes`. Captured graph layouts and the cost table must cover it; request, token and KV limits must accommodate it. | +| Verifier budget `V` | Total verifier rows in a captured rank-local compact layout: `G <= V < G*(K+1)`. V0 is a sentinel for native full-width `G*(K+1)` verification, not zero tokens. | +| Helper global `BS` | `--batch-size` is a fixed global request cohort, not a graph setting. With attention DP, G = BS / attention-DP size; otherwise G = BS. | +| Ordinary benchmark load | With attention DP, `--max_batch_size` is rank-local. `--concurrency` limits global simultaneous requests; `--num_requests` caps requests read from the dataset, with 0 meaning the dataset length. `--warmup` counts warm-up requests, not iterations or helper cohorts. | +| Model and data | Checkpoint and workload paths, using the checkpoint's tokenizer. Keep topology, precision, backend, sampling and output work fixed when comparing runs. | +| Artifacts | A compatible cost table is required for dynamic scheduling. Temperature calibration is optional. | + +The examples use K3, G64 and eight attention-DP ranks. These are example inputs, +not universal requirements. See [supported routes and configuration](CONFIDENCE_VALIDATION.md#supported-routes-and-configuration) +for runtime and helper limits. + +Observations can contain prompt and output tokens. Use a new private writable +output directory; keep inputs read-only if desired. Offline/helper commands +require fresh output paths with existing parent directories and do not overwrite +them. + +## 2. Reuse artifacts in normal runtime YAML + +Start with a complete working DeepSeek-V4 **runtime YAML**. Keep its topology, +precision, KV, backend and capacity settings, and merge this example into it, +replacing its existing speculative configuration: + +```yaml +enable_attention_dp: true +cuda_graph_config: + enable_padding: true + batch_sizes: [64] +speculative_config: + decoding_type: DSpark + max_draft_len: 3 + enable_confidence_scheduling: true + confidence_sps_table_path: /path/to/artifacts/costs.json +``` + +Save the merged runtime configuration as `dynamic.yaml`. This fragment is not a +complete configuration and does not replace the checkpoint's `config.json`. + +Create `static.yaml` from the same runtime configuration, set +`enable_confidence_scheduling: false`, and omit all confidence-only artifact, +freshness, mixed-evidence and fused-scheduler fields. Keep the physical K and +numerical/runtime settings unchanged. + +The cost table must contain finite positive native V0 measurements and the +compact cells it offers, compatible with physical K and graph/layout coverage. +Default reuse permits a configured subset of a measured graph ladder; any priced +drain tail must still be measured and captured. An enabled policy may legally +choose native throughout, but missing/incompatible coverage is not a compact +execution or a reproduced gain. + +Structurally valid artifacts can be reused across source revisions and nodes +with warnings that freshness and performance economics are unchecked. Original +artifact provenance is retained; a different destination can change prediction +quality and performance. New profiling is optional when reusing compatible +artifacts. Strict freshness and economic admission are an +[optional advanced mode](CONFIDENCE_VALIDATION.md#optional-strict-freshness-and-admission). + +## 3. Run ordinary inference or benchmarks + +Use `dynamic.yaml` through the normal `LLM` configuration or +`trtllm-serve --config` interface in [Serving](README.md#serving). + +For benchmarking, follow [DeepSeek-V4 Benchmarking](README.md#benchmarking) and +the [dataset preparation guide](../../../../docs/source/developer-guide/perf-benchmarking.md#preparing-a-dataset). +The standard benchmark JSONL format uses `task_id`, `input_ids` or `prompt`, and +`output_tokens`. Prepare at least as many requests as you intend to measure; +`--num_requests` does not repeat a shorter dataset. + +For a greedy, fixed-output comparison, save this as `sampler.yaml`: + +```yaml +temperature: 0.0 +top_k: 1 +ignore_eos: true +``` + +Run each arm in a separate process/model lifetime, changing only the config and +report paths: + +```bash +trtllm-bench --model --model_path /path/to/checkpoint throughput \ + --tp 8 --ep 8 --dataset /path/to/benchmark.jsonl \ + --max_batch_size 64 --max_num_tokens 8192 --max_seq_len 4096 \ + --concurrency 512 --num_requests 512 --warmup 512 \ + --sampler_options /path/to/sampler.yaml --config /path/to/static.yaml \ + --report_json /path/to/results/static.json + +trtllm-bench --model --model_path /path/to/checkpoint throughput \ + --tp 8 --ep 8 --dataset /path/to/benchmark.jsonl \ + --max_batch_size 64 --max_num_tokens 8192 --max_seq_len 4096 \ + --concurrency 512 --num_requests 512 --warmup 512 \ + --sampler_options /path/to/sampler.yaml --config /path/to/dynamic.yaml \ + --report_json /path/to/results/dynamic.json +``` + +Both runtime files use attention DP for this example. Match dataset output +lengths and ensure each prompt plus output fits the sequence limit. Explicit CLI +settings override conflicting YAML values. Use the normal allocation/container +procedure; `trtllm-bench` handles process launch, so do not add an MPI wrapper. + +Standard benchmark reports use the benchmark's load and timing protocol; they do +not by themselves establish compact routing. The optional helper below runs its +own routing-checked lifetimes with a different timing boundary. See the +[validation reference](CONFIDENCE_VALIDATION.md#validation-scope) for the tested +scope and interpretation of results. + +## 4. Generate artifacts when needed + +You can generate costs, temperature calibration, or both. Neither requires +generating the other first. + +### Whole-step cost table + +Follow the complete [cost-generation recipe](COST_PROFILING.md): configure a +physical-K-specific native/compact grid, profile complete executor steps, reduce +all-rank observations, and export a schema-v2 runtime table. Graph-replay timing +is a separate diagnostic, not a serving price. The +[cost format/export reference](COST_EXPORT.md) describes structural reuse checks. + +### Optional temperature calibration + +Without calibration, the head uses raw sigmoid confidence probabilities +(unit temperatures). Omit both `confidence_sts_path` in runtime YAML and +`--calibration` in helper commands to select this behavior. Omitting only the +CLI option does not remove a calibration path already present in YAML. + +To generate temperatures, use the paired collector followed by the CPU fitter. +The offline/helper input format is separate from the benchmark format: each +JSONL line contains a nonempty integer `prompt_token_ids` list. + +Set `MODEL_DIR`, `SOURCE_ROOT`, `K`, `RANKS`, `BS` and output work to your chosen +inputs. `BASELINE_YAML` must be a complete matching, confidence-OFF runtime YAML +with no confidence artifact/freshness/mixed/fused fields, explicit topology and +capacity, padded graphs, and overlap scheduling enabled. `RUN_ROOT` is a new +writable directory; `FULL_BUDGET_SECONDS` covers startup, model work and shutdown. + +```bash +mpirun -n "$RANKS" python tensorrt_llm/tools/dspark_reproduction.py run \ + --mode collect --model "$MODEL_DIR" --config "$BASELINE_YAML" \ + --inputs "$INPUTS_JSONL" --source-root "$SOURCE_ROOT" \ + --physical-k "$K" --batch-size "$BS" --output-tokens "$OUTPUT_TOKENS" \ + --timeout "$FULL_BUDGET_SECONDS" --output "$RUN_ROOT/calibration-observations" + +python tensorrt_llm/tools/dspark_calibration.py \ + --shards "$RUN_ROOT"/calibration-observations/rank-*-pairs.pt \ + --physical-k "$K" --output "$RUN_ROOT/calibration.json" +``` + +Collection runs full-K verification and pairs confidence logits with the later +accepted-prefix labels for the same proposal/request lifetime. The fitter needs +at least 1,000 genuine paired rows by default and produces per-position +temperatures, input hashes and before/after ECE. It is not head training or +task-accuracy evaluation. The [pairing contract](CONFIDENCE_VALIDATION.md#paired-collection-and-developer-checks) +describes the observations. + +To use the result, add `confidence_sts_path: /path/to/calibration.json` under +`speculative_config` in the dynamic runtime YAML, or add +`--calibration "$RUN_ROOT/calibration.json"` to a dynamic helper command. + +## 5. Optional fixed-cohort comparison and routing checks + +`dspark_reproduction.py` is for paired collection and controlled experiments, +not a required inference launcher. Its `compare` command consumes its own +`rank-N-run.json` records, not `trtllm-bench` reports. + +Use the clean `BASELINE_YAML` and offline input format described above. The helper +requires TP equal to the MPI world size, PP=CP=1 and enough rank-local capacity +for the global cohort. With attention DP, BS must divide evenly across ranks. +Run independent static and dynamic lifetimes with the same inputs, output work, +warm-up count and full budget: + +```bash +mpirun -n "$RANKS" python tensorrt_llm/tools/dspark_reproduction.py run \ + --mode static --model "$MODEL_DIR" --config "$BASELINE_YAML" \ + --inputs "$INPUTS_JSONL" --source-root "$SOURCE_ROOT" \ + --physical-k "$K" --batch-size "$BS" --output-tokens "$OUTPUT_TOKENS" \ + --warmup-cohorts 1 --timeout "$FULL_BUDGET_SECONDS" --output "$RUN_ROOT/static" + +mpirun -n "$RANKS" python tensorrt_llm/tools/dspark_reproduction.py run \ + --mode dynamic --model "$MODEL_DIR" --config "$BASELINE_YAML" \ + --inputs "$INPUTS_JSONL" --source-root "$SOURCE_ROOT" \ + --physical-k "$K" --batch-size "$BS" --output-tokens "$OUTPUT_TOKENS" \ + --cost-table "$COST_TABLE_JSON" \ + --warmup-cohorts 1 --timeout "$FULL_BUDGET_SECONDS" --output "$RUN_ROOT/dynamic" + +python tensorrt_llm/tools/dspark_reproduction.py compare \ + --runs "$RUN_ROOT/static/rank-0-run.json" "$RUN_ROOT/dynamic/rank-0-run.json" \ + --output "$RUN_ROOT/comparison.json" +``` + +The helper records resolved/runtime enablement, native/compact replay, outputs, +source/artifact identity and shutdown. `compare` requires actual compact routing +for a compact-gain claim; an enabled all-native run is not reported as one. +Its TPS covers fixed-cohort enqueue through final response, including prefill +and drain. Use balanced fresh-lifetime repetitions and equal static/dynamic +record counts. Do not pool these percentages with standard benchmark or plateau +measurements. Token equality is diagnostic, not a task-accuracy test. + +Advanced configuration and optional developer checks are in the +[validation reference](CONFIDENCE_VALIDATION.md). diff --git a/examples/models/core/deepseek_v4/CONFIDENCE_VALIDATION.md b/examples/models/core/deepseek_v4/CONFIDENCE_VALIDATION.md new file mode 100644 index 000000000000..affda300eb3a --- /dev/null +++ b/examples/models/core/deepseek_v4/CONFIDENCE_VALIDATION.md @@ -0,0 +1,260 @@ + + + +# Confidence scheduling: validation and developer reference + +Start with the [usage guide](CONFIDENCE_REPRODUCTION.md) for artifact reuse, +self-generation and ordinary inference/benchmarking. This reference records +supported/tested scope and advanced contracts. Developer checks and strict +validation are opt-in, not an ordinary pre-run checklist. + +## Validation scope + +These are the distinct integration scopes recorded before this documentation +reorganization. Original source/native/artifact provenance remains unchanged. + +| Scope | Recorded result | Does not establish | +| --- | --- | --- | +| Earlier physical-K3, global-BS512/G64 across eight attention-DP ranks, separately recorded Python/native environment | Full-model whole-step profiling/export, genuinely paired calibration collection/fitting, and independent static/enabled-compact generation completed with output and natural-shutdown checks. | Fresh measurements for another source/runtime/node or every other K/BS. | +| Newer source composition, same-source/same-node K3, global-BS512/G64, 512 requests × 32 output tokens | Native/API compatibility and independent static/enabled-compact lifetimes completed source/native/artifact/input/output/routing, public comparison and natural-shutdown checks, reusing original costs/calibration. Exact generated token sequences differed. | Exact-source native rebuild, global clean install, task accuracy or sustained-serving performance. | +| Ordinary `trtllm-bench --config` confidence recipe in the usage guide | Public CLI/YAML/sampler interfaces are source-supported. | A new GPU execution or performance result for that recipe. | +| Small-K drafter-fusion extension | Source-matched extension passed 98 GPU unit cases on B300/SM103. | Full-model/performance qualification of every composed integration, SM100 qualification or arbitrary-K kernel support. | + +Other prepared batch sizes and configurable physical K are interface coverage, +not additional measured results. Existing source tests and historical +measurements must not be relabeled as current-head GPU validation. +The interfaces support the global BS128/256/512/1024/1536/2048 grid and smaller +points when actual graph/K/cost/capacity coverage permits. With eight attention-DP +ranks these map to G16/G32/G64/G128/G192/G256; with attention DP off, G is the +global cohort. This mapping is not a GPU-validation claim for every point. + +The controlled helper reports fixed-cohort enqueue-to-final-response TPS, +including prefill and drain. Historical plateau measurements and ordinary +benchmark load protocols use different timing boundaries; percentages from them +must not be pooled. Exact output-token equality is a diagnostic field, not a +replacement for original task-accuracy tests or a universal bitwise requirement. +No confidence interval or best-static/best-dynamic claim is produced by one pair. + +## Supported routes and configuration + +The confidence route currently supports embedded DeepSeek-V4 DSpark with trained +confidence-head weights, PP=1 and padded CUDA graphs. Standalone drafters retain +native verification; disable confidence scheduling for that route. The ordinary +target verification, sampling/acceptance, KV/graph, kernel and fatal-error checks +remain in force. + +Guided decoding, LoRA, per-request logits post-processors and incompatible sparse +top-k paths are unsupported. In particular, `enable_heuristic_topk` and +`use_cute_dsl_topk` assume uniform speculative strides. On ragged steps, +`use_cute_dsl_paged_mqa_logits` falls back to the expanded DeepGEMM path. +Advanced rejection draft sampling is unsupported because its sampled prefix +differs from the greedy prefix that produced the confidence scores. Greedy +requests, including `top_k=1`, and requests with rejection draft sampling +disabled retain their supported behavior. Confidence-off execution retains its +ordinary sampling/post-processing paths. + +The offline collector and fixed-cohort helper additionally require aggregate +PP=1/CP=1 generation, embedded DSpark, MPI world matching TP, sufficient +rank-local request/token/KV capacity and padded decode graphs. They reject +guided decoding, LoRA, custom logits processing, disaggregated KV transfer, +in-graph sampling, encode-only and resource-governor control. Paired collection +uses the original overlap-scheduler seam. These helper restrictions are not a +claim that all ordinary TensorRT LLM routes share the same limitations. + +### Cost/configuration contract + +`V=0` denotes native full-width verification; positive V denotes a measured +compact layout at rank-local graph size G. Physical K is not reduced by +scheduling. K1 has no compact tier and resolves to feature-off static execution. +Equal row counts with different graph/attention layouts are not interchangeable. + +The schema-v2 table must contain finite positive measurements and compatible +physical-K, graph/layout and capacity coverage. Default reuse can select a +configured subset of the original measured ladder while validating the entire +original table/fingerprint; a drain-guard tail must still be measured and +captured. Strict freshness requires the exact original ladder. Missing cells +are not interpolated, copied from another K or fabricated. + +The current whole-step exporter supports attention-DP fanout or one rank; +multi-rank attention-DP-off export is refused. This is an export-schema limit, +not a blanket DSpark runtime restriction. Replay-entry timing remains a separate +diagnostic; it is not renamed to a complete-loop cost or summed with overlapping +host/GPU spans. + +Temperature vectors, when supplied, must have finite positive values for each +physical draft position. See the [optional calibration workflow](CONFIDENCE_REPRODUCTION.md#optional-temperature-calibration) +for collection, fitting and use. Reuse retains original artifacts and records +the actual destination runtime separately; it does not rewrite old fingerprints +or guarantee positive performance. + +### Requested configuration versus execution + +| Case | Resolved behavior | +| --- | --- | +| Explicit confidence enablement, default artifact reuse | Enabled, subject to structural/runtime checks; warnings state freshness/economics are unchecked. | +| Optional strict freshness with missing, stale, unreadable or non-positive evidence | Feature-off static; confidence flags and artifact paths/digests are cleared with a warning. | +| Physical K1 | Feature-off static K1, confidence inputs cleared. | +| Strict freshness with matching positive evidence | Remains requested; normal runtime/checkpoint/graph checks still apply. | + +Startup feature-off fallback is not an enabled-native policy decision. An enabled +policy can choose native while retaining its confidence/selection work. Optional +fused-scheduler failure can instead fall back to the tensor scheduler for an +affected graph size. Record resolved and runtime enablement plus actual routing; +a YAML flag or successful server start is insufficient. + +The selector retains the table's immediate predicted-goodput margin, normally +1%, and any supplied iteration/drain guard. Default reuse without a drain guard +uses immediate goodput; strict validation without one retains native +verification. Fixed confidence-path overhead is a separate workload-level +charge, not redistributed to individual candidate choices. None certifies an +end-to-end speedup. + +### Drafter fusion + +The BF16 CuTe attention and RMSNorm–RoPE path is eligible for physical K=1..8 +under the existing SM100/SM103 geometry/layout guards; unsupported inputs retain +the reference path. The eight-row storage bound is unchanged. This dispatch is +independent of the confidence flag and therefore also affects static DSpark. +See the validation table above for the actual exercised scope. + +## Optional strict freshness and admission + +`validate_confidence_artifact_freshness: true` opts into the stricter path. +It consumes four externally supplied inputs: + +- `confidence_sps_table_path`: schema-v2 measured costs. +- `confidence_sps_live_fingerprint_path`: independently observed destination + fingerprint, not a relabeled historical fingerprint. +- `confidence_admission_receipt_path`: schema-v1 workload receipt. +- `confidence_admission_receipt_sha256`: externally trusted SHA256 of its bytes. + +Ordinary artifact reuse does not require these additional artifacts. The public +cost/calibration commands do not produce a positive strict-admission receipt. +A separately reviewed producer must retain measured costs, policy overhead, +calibration, workload, selector-replay and runtime evidence. Source, workload or +artifact changes require reassessing that evidence's applicability. + +The receipt envelope has exactly `schema_version: 1`, `admission` and +`admission_sha256`. The last field hashes the canonical admission body: +`json.dumps(body, sort_keys=True, separators=(",", ":")).encode()`. +Its fields are: + +```text +physical_k, source_head, source_diff_sha256, runtime_snapshot +sps_cost_table_sha256, live_engine_fingerprint_sha256, selector_identity_sha256 +workload_identity_sha256, calibration_result_sha256, selector_replay_sha256 +policy_steps, compact_choices, gross_compact_value_ms_lower_bound +fixed_confidence_path_overhead_ms_upper_bound, safety_margin_ms, admitted +``` + +The selector identity must match `EXACT_SPS_SELECTOR_IDENTITY_SHA256` in +`tensorrt_llm/_torch/speculative/dspark_planner.py`. Admission recomputes: + +```text +net_value_lower_bound_ms = + gross_compact_value_ms_lower_bound + - policy_steps * fixed_confidence_path_overhead_ms_upper_bound + - safety_margin_ms +admitted = compact_choices > 0 and net_value_lower_bound_ms > 0 +``` + +All policy steps, including native choices, pay enabled-path overhead. Do not +allocate that fixed charge back to candidates using a +`policy_steps / compact_choices` ratio. Admission does not replace per-decision +selection or its original drain/price guards. + +Hashes check integrity/cross-artifact consistency, not signatures or origin. +The consumer does not independently derive running checkpoint/workload identity, +rerun calibration/replay or verify every referenced result file. The external +producer remains the trust boundary; hashing arbitrary JSON or fixtures is not +admission evidence. + +## Optional mixed context/generation scheduling + +`confidence_mixed_evidence` defaults to `None`: mixed steps use native +verification while supported decode-only scheduling remains enabled. Decode +T(G,V) costs do not price mixed work; context chunks/positions, history, +bootstrap rows, ordered ranks and shared body capacity enter the mixed key. + +The evidence object has `costs_path`, `costs_sha256`, `admission_path`, +`admission_sha256`, `protocol_sha256`, `live_fingerprint_path` and +`live_fingerprint_sha256`. Cost and receipt schemas are +`dspark-mixed-breakable-costs-v1` and `dspark-mixed-breakable-admission-v1`. +Pins come from a trusted producer; the live source/runtime/model/GPU/topology/ +selector fingerprint must be observed independently. Decode artifact reuse does +not relax this separate contract. + +The supported mixed route is embedded physical-K3 DSpark, TP8/PP1/CP1, +attention DP, greedy sampling and captured BREAKABLE buckets +`[96, 128, 192, 256, 384, 512]`. Unsupported/uncovered routes stay native for +mixed steps without disabling valid decode scheduling. First context chunks +stay native because prefix-cache resolution can change geometry. Later chunks +need their full-width native rows to fit; quoted context work is not trimmed. + +When a peer prefills, ranks use shared BREAKABLE padding, not a rank-local full +decode graph. Only authenticated generation owners can receive shorter +prefixes; context/bootstrap/padding rows retain required rows. Host-authoritative +lengths are published before attention metadata preparation, without changing +backend, sampler or acceptance. Stale/unready owners remain full-width. + +Selection requires an exact measured geometry/action cell and at least 1% +predicted-goodput margin. Current `rank0_forward_pricing_proxy` prices forward +work, not the whole distributed step. Workload admission must charge CPU curve +preparation, producer-stamp transfer and the existing allgather's fixed trailer +on all policy steps, including native/all-decode decisions. Policy equivalence +does not establish zero overhead. + +There is no public mixed receipt-generation command here. Valid-shaped fixture +JSON is not deployment evidence. A separately reviewed producer must supply +measured costs, accuracy/runtime results, selector replay and positive workload +economics; exact integration GPU/full-model and production evidence remain +separate from CPU contracts. + +## Paired collection and developer checks + +Calibration pairs are not matched by arrival order. Step n confidence describes +a proposal verified later, usually at n+1. The adapter binds that forward to its +original `SampleState.device`, snapshots proposals before buffer reuse, and +records confidence/draft stamps before the normal sampler event. Host copies +are read only after the sampler's existing synchronization/update. + +Request incarnation, proposal content, local confidence row, sampler slot and +full-K extent must agree. Prefill placeholders, retired/unscored bootstrap rows +are counted rather than fitted. Wrong-step/stale/censored observations fail; +current-step logits cannot substitute for the producing logits. Compact +verification cannot label an unverified suffix. + +Collection emits `rank-N-pairs.pt` and `rank-N-run.json`. Tensor-only shards +contain logits [rows,K], cumulative accepted-prefix labels, pairing metadata and +source/runtime identity. The fitter retains shard hashes, sample count, +temperatures and before/after ECE; degenerate positions are reported rather than +promoted as informative fits. This is not task-accuracy qualification. + +Artifact, calibration and observation controls can run with CPU inputs: + +```bash +python -B tests/unittest/tools/test_dspark_cost_export.py +python -B tests/unittest/tools/test_dspark_calibration.py +python -B tests/unittest/tools/test_dspark_observations.py +``` + +The constructor test uses the real argument validator and internal factory seam +without constructing a model. Run it in a compatible installed native runtime; +imports may require a driver: + +```bash +python -m pytest -q tests/unittest/executor/test_dspark_reproduction_factory.py +``` + +`executor_cls` is internal and remains rejected as a public LLM argument. +The factory tests cover observation/profile construction, the imported replay +keyword/default contract, position-offset forwarding and routing bookkeeping. +These CPU/native-import controls are not real MPI, GPU execution, accuracy or +performance results. Additional contract regressions are: + +- `tests/unittest/_torch/speculative/hw_agnostic/test_dspark_workload_admission.py` +- `tests/unittest/_torch/speculative/hw_agnostic/test_dspark_admission_resolution.py` +- `tests/unittest/_torch/speculative/hw_agnostic/test_dspark_artifact_reuse.py` +- `tests/unittest/llmapi/test_dspark_workload_admission_config.py` +- `tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_evidence.py` +- `tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_runtime.py` diff --git a/examples/models/core/deepseek_v4/COST_EXPORT.md b/examples/models/core/deepseek_v4/COST_EXPORT.md new file mode 100644 index 000000000000..290a0d183fb2 --- /dev/null +++ b/examples/models/core/deepseek_v4/COST_EXPORT.md @@ -0,0 +1,72 @@ + + + +# Cost format, export and reuse + +`tensorrt_llm/tools/dspark_cost_export.py` accepts either a reduced whole-step +summary or an existing schema-v2 runtime cost table. It validates the requested +physical K and graph ladder using the runtime's structural validator. + +For the complete GPU profile → CPU reduce → CPU export procedure, use +[Generate DSpark whole-step costs](COST_PROFILING.md). For normal runtime +configuration, start with the [reproduction guide](CONFIDENCE_REPRODUCTION.md). + +## Accepted timing scopes + +| Input | Treatment | +| --- | --- | +| Whole-step `verify-cost-measured-summary-v1` | Requires `original-executor-loop-hostStepTimeMS` from the explicit whole-step producer. Direct retained medians become schema-v2 `step_time_ms`. | +| Existing schema-v2 runtime table | Structurally validated without changing original bytes, measurements or fingerprint. | +| Graph-replay summary | Rejected as a runtime-price source. Replay-entry medians cannot be relabeled whole-step costs. | + +Whole-step observations use the executor's existing complete-loop timer and +exact all-rank iteration joins. The adapter takes the median of retained +per-iteration rank maxima without summing overlapping graph/host spans. Each +declared native/compact cell requires its own direct measurement, with at least +12 retained samples per generation and 24 pooled per cell. + +The exported fingerprint and source/configuration observations describe the +measurement's original execution. Export does not reobserve the destination or +turn loaded-module hashes into a newly measured live fingerprint. + +## Check or copy an existing runtime table + +For example, to use a table with physical K3 and a runtime capturing G16 and G64: + +```bash +python tensorrt_llm/tools/dspark_cost_export.py \ + --input /path/to/costs.json --physical-k 3 --graph-batch-sizes 16 64 +``` + +Pass the actual resolved rank-local graph ladder. Default reuse permits a +configured subset of the measured ladder while validating the entire original +table and fingerprint. Native V0 coverage, offered compact layouts, physical K, +finite positive timings and cell provenance must remain compatible. Any priced +drain tail must be both measured and captured. Missing prices are not +interpolated, copied from another K or filled with zero. + +Add `--output /path/to/new-costs.json` for a byte-identical copy of a valid +schema-v2 input. The destination must not exist. Without `--output`, the command +only prints its structural-check report. + +Structural validation does not require matching source HEAD/node or economic +admission. It warns that destination freshness and performance are unchecked. +Confidence remains default-OFF and needs explicit enablement in runtime YAML. +The separately optional [strict freshness/admission mode](CONFIDENCE_VALIDATION.md#optional-strict-freshness-and-admission) +has additional evidence requirements, including the exact original ladder. + +## New whole-step tables + +When the input is a reduced whole-step summary, `--output` writes a new runtime +table rather than copying summary bytes. Every measurement references the +original summary hash. The default `--minimum-predicted-gain 0.01` sets a policy +margin; it is not a timing measurement. + +The current schema checks global graph counts as rank-local G times GPU count. +New whole-step export therefore supports attention-DP fanout or one rank and +explicitly refuses multi-rank attention-DP-off export. This is an export-schema +limit, not a general DSpark runtime limitation. + +This tool does not fit temperatures, measure policy overhead, certify task +accuracy, grant admission or prove a speedup. Those checks remain separate from +the timing-format conversion. diff --git a/examples/models/core/deepseek_v4/COST_PROFILING.md b/examples/models/core/deepseek_v4/COST_PROFILING.md new file mode 100644 index 000000000000..a4c6d0bb636e --- /dev/null +++ b/examples/models/core/deepseek_v4/COST_PROFILING.md @@ -0,0 +1,172 @@ + + + +# Generate DSpark whole-step costs + +Use this recipe when you need a new cost table. Reusing a compatible supplied +table does not require profiling again. Installation, artifact reuse and +ordinary execution are covered by the [reproduction guide](CONFIDENCE_REPRODUCTION.md). + +The workflow is **profile → reduce → export**. Whole-step collection uses the +executor's complete-step `hostStepTimeMS` timer rather than only graph replay. +It uses the normal model, KV, sampler and +shutdown paths. No forced acceptance is permitted. + +## 1. Configure the runtime and grid + +K is the physical draft block, G is the rank-local request count, and V is the +total verifier-row budget of a captured compact layout. V0 is the sentinel for +native full-width `G*(K+1)` verification, not zero work. Every G needs a native +V0 control. Compact cells require K > 1 and `G <= V < G*(K+1)`; choose the compact +cells you want to price rather than sweeping every possible V. + +Prepare a complete runtime YAML for the intended model, precision and backend. +It must have: + +- `tensor_parallel_size` equal to the MPI world size, PP=1 and CP=1; +- the intended expert parallelism and attention DP, plus sufficient request, + token, sequence-length and KV capacity; +- `enable_iter_perf_stats: true` for the complete-loop timer; +- padded CUDA graphs with a resolved rank-local ladder exactly matching all G + values in the grid; +- embedded DSpark at the chosen physical K, with confidence scheduling and the + fused scheduler OFF, no mixed evidence or draft-length schedule, and no + SPS, STS, fingerprint or admission fields. + +Do not use an enabled `dynamic.yaml` unchanged for profiling. The grid's +`record_confidence_head: true` records the loaded head inside the graphs without +enabling serving confidence policy. + +The current whole-step exporter supports attention-DP fanout or a single rank; +multi-rank attention-DP-off export is refused. The offline driver supports +aggregate greedy generation and rejects guided decoding, LoRA, custom logits +processors, disaggregated KV transfer/connectors, in-graph sampling, encode-only +and resource-governor control. Multi-rank operation needs MPI `THREAD_MULTIPLE`. + +For example, use this `profile-grid.json` for K3/G64: + +```json +{ + "physical_k": 3, + "record_confidence_head": true, + "max_samples": 4096, + "cells": [ + {"graph_batch_size": 64, "verifier_budget": 0}, + {"graph_batch_size": 64, "verifier_budget": 96}, + {"graph_batch_size": 64, "verifier_budget": 128}, + {"graph_batch_size": 64, "verifier_budget": 192} + ] +} +``` + +This declares shapes, not measured prices. Merge the following fragment into +the complete working runtime YAML for this example; it is not a standalone +configuration: + +```yaml +tensor_parallel_size: 8 +pipeline_parallel_size: 1 +context_parallel_size: 1 +moe_expert_parallel_size: 8 +enable_attention_dp: true +enable_iter_perf_stats: true +cuda_graph_config: + enable_padding: true + batch_sizes: [64] +speculative_config: + decoding_type: DSpark + max_draft_len: 3 + enable_confidence_scheduling: false +``` + +Change K and the grid/graph ladder together for another workload. Each additional +G needs its own native control and desired compact cells. Request/token limits +must retain the requested ladder: a silently clamped graph is not its requested +shape. `max_samples` must accommodate the run's raw replay records. + +The JSONL workload must contain at least the largest global cohort, with a +nonempty integer `prompt_token_ids` list per line. Use the model's tokenizer and +ensure prompt plus output fits `max_seq_len`. Each cell repeats the same real +prompt cohort through prefill, generation and drain. Only decode steps with +exactly G requests on every rank and the declared graph key are samples; +bootstrap and shorter drain steps retain native execution. + +## 2. Profile complete steps on GPUs + +Run from the source checkout using its compatible installed runtime and the +normal allocation/container/MPI procedure. Set these Bash variables to your +paths and selected inputs; for the example, `K=3`, `RANKS=8` and `GRAPH_G=(64)`. +`SOURCE_ROOT` is the checkout used for execution; `PROFILE_YAML` is the complete +confidence-OFF configuration above. Output parents must exist and output paths +must be new. + +```bash +mpirun -n "$RANKS" python -m tensorrt_llm.tools.verify_costs profile \ + --model "$MODEL_DIR" --config "$PROFILE_YAML" --grid "$GRID_JSON" \ + --physical-k "$K" --cost-scope whole-step --source-root "$SOURCE_ROOT" \ + --inputs "$INPUTS_JSONL" --output "$PROFILE_OUT" \ + --generations 3 --output-tokens 768 --timeout "$GENERATION_BUDGET_SECONDS" +``` + +`--timeout` is the shared generation-phase budget after construction and identity +checks, not the whole allocation budget. Reserve startup, generation, export and +natural shutdown in the outer launcher. Keep the full chosen output workload; +a short run may not meet the sample floors below. + +The output contains `rank-R-record-N.json` for each cell/generation/rank, with +actual graph keys, per-step samples, request/sampler/KV joins, responses and +source/runtime/model/topology identity. Only completed runs with successful +shutdown can be exported as complete measurements. + +## 3. Reduce all-rank observations on CPU + +This reduction uses the standard library and does not import native TensorRT LLM: + +```bash +python tensorrt_llm/tools/verify_costs.py export \ + --raw "$PROFILE_OUT"/rank-*-record-*.json --output "$SUMMARY_JSON" \ + --exclude-generations 1 --warmup-samples 2 \ + --minimum-generation-samples 12 --minimum-samples 24 +``` + +The shell expands the raw-file glob. Generation 0 is excluded, as are the first +two eligible samples of each retained generation. Every retained generation +needs at least 12 exact samples, and every cell needs at least 24 pooled samples. +More samples in one generation do not waive another generation's floor. + +The reducer joins replay ordinals and executor iterations across all ranks, +takes each joined rank maximum, then reports the retained median. It does not +average ranks or sum overlapping spans. It rejects missing/duplicate ranks or +cells, changed layouts, unequal identities, nonfinite/nonpositive times and +incomplete closure. Each cell must retain one consistent observed graph key; +different sequence-length/attention graph regimes need separate declared runs. + +## 4. Export the runtime table on CPU + +```bash +python tensorrt_llm/tools/dspark_cost_export.py \ + --input "$SUMMARY_JSON" --physical-k "$K" \ + --graph-batch-sizes "${GRAPH_G[@]}" --output "$COST_TABLE_JSON" +``` + +`GRAPH_G` is a Bash array of the resolved ladder, for example `(16 64)` for two +sizes. The exporter maps directly measured whole-step medians to schema-v2 +`step_time_ms`, retaining their original source/fingerprint and summary hash. +It does not invent missing prices or refresh a historical identity. Use the +new table as `confidence_sps_table_path` in the normal runtime configuration. + +See [Cost format, export and reuse](COST_EXPORT.md) for accepted formats and +structural checks. Costs do not calibrate the head, establish task accuracy or +grant economic admission. + +## Advanced: graph-replay diagnostics + +The same `profile` CLI defaults to `--cost-scope graph-replay`. That mode measures +graph replay including static input copies. Use it for +graph-level investigation, omitting the whole-step-only `--source-root` option. +Its replay medians are not serving prices and are rejected by the schema-v2 cost +exporter. GPU-forward timing and overlap `iterLatencyMS` are not substitutes for +the complete-loop boundary, and overlapping host/GPU spans must not be summed. + +Supported/tested scope and optional developer controls are in the +[validation reference](CONFIDENCE_VALIDATION.md). diff --git a/examples/models/core/deepseek_v4/README.md b/examples/models/core/deepseek_v4/README.md index 5672ef03a73e..bed6e1cca423 100644 --- a/examples/models/core/deepseek_v4/README.md +++ b/examples/models/core/deepseek_v4/README.md @@ -36,6 +36,7 @@ for how to build TensorRT LLM from source and start a TRT-LLM Docker container. - [Sparse attention](#sparse-attention) - [KV cache](#kv-cache) - [Quantized checkpoints](#quantized-checkpoints) + - [DSpark confidence verification (prototype)](#dspark-confidence-verification-prototype) - [Notes and Troubleshooting](#notes-and-troubleshooting) @@ -472,6 +473,16 @@ MXFP4 routed MoE expert weights, TensorRT LLM automatically applies the routed-e configuration. +## DSpark confidence verification (prototype) + +DSpark can use a trained confidence head to choose a measured verifier-token budget while +continuing to propose the full physical draft block. Confidence scheduling is experimental and +off by default. Enabling it requires a compatible cost table. + +Start with [Confidence scheduling and reproduction](CONFIDENCE_REPRODUCTION.md) for ordinary +runtime YAML and benchmarking, artifact reuse or self-generation, and optional temperature +calibration and controlled comparisons. + ## Notes and Troubleshooting - `DeepSeek-V4 on Hopper requires kv_cache_config.dtype='fp8_ds_mla'`: remove an explicit BF16/FP8 diff --git a/setup.py b/setup.py index e22f2bbf07a5..b3d5a7500714 100644 --- a/setup.py +++ b/setup.py @@ -182,11 +182,11 @@ def has_ext_modules(self): 'libs/libtensorrt_llm_ucx_wrapper.so', 'libs/libdecoder_attention_0.so', 'libs/libtensorrt_llm_nixl_wrapper.so', - 'libs/nixl/**/*', + 'libs/nixl/*', 'tensorrt_llm_transfer_agent_binding*.so', 'tensorrt_llm_transfer_agent_binding.pyi', 'libs/libtensorrt_llm_mooncake_wrapper.so', - 'libs/ucx/**/*', + 'libs/ucx/*', 'libs/libpg_utils.so', 'libs/libdecoder_attention_1.so', 'libs/nvshmem/License.txt', @@ -196,7 +196,7 @@ def has_ext_modules(self): 'deep_ep/LICENSE', 'deep_ep/*.py', 'deep_ep_cpp_tllm.*.so', - "include/**/*", + "include/*", 'deep_gemm/LICENSE', 'deep_gemm/include/**/*', 'deep_gemm/*.py', diff --git a/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py index d414a99f60a8..8feab7b891bb 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py @@ -5,7 +5,7 @@ The op runs the tcgen05 MMA kernel (:class:`DSparkAttention`) on the supported DSV4 DSpark -geometry: 128 heads, head_dim 512, draft block 5 or 6, and a 128-row rolling +geometry: 128 heads, head_dim 512, draft block 1 through 8, and a 128-row rolling window. Other shapes fall back to the pure-PyTorch reference path in ``models/modeling_dspark.py``. """ @@ -32,7 +32,8 @@ _DSV4_DSPARK_HEAD_DIM = 512 _DSV4_DSPARK_WINDOW_SIZE = 128 _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE = 8 -_DSV4_DSPARK_BLOCK_SIZES = (5, 6) +# Physical K is bounded by the eight-row draft-buffer ABI, not the MMA tile. +_DSV4_DSPARK_BLOCK_SIZES = (1, 2, 3, 4, 5, 6, 7, 8) _DSV4_DSPARK_ROPE_DIM = 64 @@ -92,7 +93,9 @@ def is_fused_dsv4_dspark_attention_supported( f"expected q/kv_cache ranks 4/3, got {q.ndim}/{kv_cache.ndim}", "tensor_ranks" ) if q.shape[1] not in _DSV4_DSPARK_BLOCK_SIZES: - return _log_unsupported(f"draft block size must be 5 or 6, got {q.shape[1]}", "block_size") + return _log_unsupported( + f"draft block size must be 1 through 8, got {q.shape[1]}", "block_size" + ) if q.shape[2:] != (_DSV4_DSPARK_NUM_HEADS, _DSV4_DSPARK_HEAD_DIM): return _log_unsupported( "q must have 128 heads and head_dim 512; " @@ -346,7 +349,7 @@ def fused_dsv4_dspark_attention( ): raise ValueError( "fused_dsv4_dspark_attention requires contiguous supported DSV4 DSpark tensors " - "([B, 5|6, 128, 512] queries, [B, 8, 512] draft blocks, INT32 indices, " + "([B, K, 128, 512] queries with 1 <= K <= 8, [B, 8, 512] draft blocks, INT32 indices, " "and a 128-row BF16 window) on SM100 or SM103" ) diff --git a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py index 56bc9f62d707..f2e0b741b2b5 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py @@ -20,7 +20,7 @@ _DSV4_DSPARK_ROPE_DIM = 64 _DSV4_DSPARK_WINDOW_SIZE = 128 _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE = 8 -_DSV4_DSPARK_BLOCK_SIZES = (5, 6) +_DSV4_DSPARK_BLOCK_SIZES = (1, 2, 3, 4, 5, 6, 7, 8) _DSV4_DSPARK_ARCH_BY_SM = { 100: "sm_100", 103: "sm_103", diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark/attention.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark/attention.py index 458d2877f12d..098458969918 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark/attention.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark/attention.py @@ -195,7 +195,7 @@ def __new_from_mlir_values__(self, values): class DSparkAttention(DSparkAttentionKernel): - """Attention over a 128-token rolling window and one 5/6-token draft block.""" + """Attention over a 128-token rolling window and one supported draft block.""" window_size = 128 block_size = 6 @@ -260,8 +260,9 @@ def __init__( ] if mismatches: raise ValueError("Unsupported DSpark kernel configuration: " + ", ".join(mismatches)) - if seq_len_q not in (5, 6): - raise ValueError(f"DSpark block size must be 5 or 6, got {seq_len_q}") + # Each query gets its own head tile; only draft storage bounds K. + if seq_len_q not in (1, 2, 3, 4, 5, 6, 7, 8): + raise ValueError(f"DSpark block size must be 1 through 8, got {seq_len_q}") if inverse_rope_dim not in (0, 64): raise ValueError(f"DSpark inverse_rope_dim must be 0 or 64, got {inverse_rope_dim}") diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index ccb4609d09c8..e05538bbd597 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -2103,13 +2103,23 @@ def write_context_windows_batched(self, main_hidden, positions, slots, mask, kv_ ) def _prepare_confidence_head_for_load(self, weights: Dict) -> None: - """Match the projection's optional bias to the checkpoint before loading.""" + """Prepare and validate the embedded confidence namespace before loading.""" head = self.dspark_model.mtp_layers[-1].confidence_head - bias_key = f"mtp_layers.{self.num_stages - 1}.confidence_head.proj.bias" - if head is not None and bias_key in weights and head.proj.bias is None: + prefix = f"mtp_layers.{self.num_stages - 1}.confidence_head." + module_weights = { + name[len(prefix) :]: value for name, value in weights.items() if name.startswith(prefix) + } + if head is None: + if module_weights: + raise ValueError("DSpark confidence head checkpoint has no constructed head") + return + if "proj.bias" in module_weights and head.proj.bias is None: # This runs before weight loading and graph capture. Empty storage # avoids consuming RNG or reinitializing the existing projection. head.proj.bias = nn.Parameter(head.proj.weight.new_empty(head.proj.out_features)) + # The generic V4 loader skips this parameter-free head container and + # visits its Linear child, so strict validation must run at this seam. + head._validate_checkpoint_weights(module_weights) def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): """Load the ``mtp.*`` draft weights from the (full) checkpoint dict. diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 569ea56dced9..52d3134b55d4 100644 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -579,6 +579,13 @@ def load_sts_temperatures(self, temperatures: torch.Tensor) -> None: def load_weights(self, weights: list) -> None: """Load every confidence-head checkpoint parameter strictly.""" (module_weights, ) = weights + self._validate_checkpoint_weights(module_weights) + for name, param in self.named_parameters(): + param.data.copy_(module_weights[name][:]) + + def _validate_checkpoint_weights( + self, module_weights: dict[str, torch.Tensor]) -> None: + """Validate a head namespace before generic embedded parameter loading.""" expected = dict(self.named_parameters()) unexpected = [key for key in module_weights if key not in expected] missing = [key for key in expected if key not in module_weights] @@ -593,7 +600,11 @@ def load_weights(self, weights: list) -> None: f"DSpark confidence head checkpoint is missing key(s) {sorted(missing)}" ) for name, param in expected.items(): - param.data.copy_(module_weights[name][:]) + actual_shape = tuple(module_weights[name].shape) + if actual_shape != tuple(param.shape): + raise ValueError( + f"DSpark confidence head checkpoint {name} has shape {actual_shape}; " + f"expected {tuple(param.shape)}") class Eagle3Attention(Attention): diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 2fde7afd6d5e..bab0d8dbbcba 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -85,12 +85,23 @@ from .seq_slot_manager import SeqSlotManager GB = 1 << 30 +_DEFAULT_INDEX_MAPPER_RESERVED_SLOTS = 1 +_CONFIDENCE_INDEX_MAPPER_RESERVED_SLOTS = 2 def ceil_div(a: int, b: int) -> int: return (a + b - 1) // b +def _get_num_reserved_index_slots(spec_config: Optional[SpeculativeConfig], + estimating_kv_cache: bool) -> int: + """Return IndexMapper headroom required by the execution policy.""" + if getattr(spec_config, "enable_confidence_scheduling", False): + return (0 if estimating_kv_cache else + _CONFIDENCE_INDEX_MAPPER_RESERVED_SLOTS) + return _DEFAULT_INDEX_MAPPER_RESERVED_SLOTS + + def _get_initial_lora_data_type( configured_lora_data_type: Optional[torch.dtype], ) -> Optional[torch.dtype]: @@ -3278,6 +3289,8 @@ def _create_kv_cache_manager( manager_extra_kwargs["max_copy_beam_width"] = max_copy_beam_width manager_extra_kwargs[ "cold_page_codec_provider"] = cold_page_codec_provider + manager_extra_kwargs["num_reserved_index_slots"] = ( + _get_num_reserved_index_slots(spec_config, estimating_kv_cache)) manager_extra_kwargs["kv_events_config"] = kv_events_config manager_extra_kwargs["joint_kv_cache_reuse"] = joint_kv_cache_reuse manager_extra_kwargs[ @@ -3984,6 +3997,7 @@ def create_py_executor_instance( execution_stream: Optional[torch.cuda.Stream] = None, dwdp_manager: Optional[DwdpManager] = None, max_num_sequences: Optional[int] = None, + executor_type: Optional[type[PyExecutor]] = None, ) -> PyExecutor: set_low_latency_dispatch( getattr(llm_args, 'enable_low_latency_host_dispatch', False)) @@ -4353,7 +4367,8 @@ def create_py_executor_instance( executor_max_seq_len = min(executor_max_seq_len, int(decoder_position_limit)) - return PyExecutor( + executor_type = PyExecutor if executor_type is None else executor_type + return executor_type( resource_manager, scheduler, model_engine=model_engine, diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 19ed990f33a1..6b55382bb149 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( @@ -281,6 +349,11 @@ def _is_mixed_encoder_decoder_batch(self, batch: ScheduledRequests) -> bool: and batch.num_generation_requests > 0) def _can_run_cuda_graph_batch(self, batch: ScheduledRequests) -> bool: + # A peer may be prefilling while this rank contains only generation. + # Host-authoritative windows use the common BREAKABLE body, never a + # rank-local full-generation graph or its device layout prologue. + if getattr(self, "_dspark_host_window_batch", None) is batch: + return False return (batch.can_run_cuda_graph or (self.enable_encoder_decoder_mixed_cuda_graph and self._is_mixed_encoder_decoder_batch(batch))) @@ -430,7 +503,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 +581,32 @@ def _get_compatible_mixed_encoder_decoder_key(self, return min(compatible_keys, key=lambda captured_key: captured_key.num_encoder_tokens) + def _ragged_verify_bucket(self, batch: ScheduledRequests) -> Optional[int]: + """Optional token bucket supplied by a specialized graph runner.""" + return None + + def _agree_graph_shape( + self, batch: ScheduledRequests, can_run_cuda_graph: bool, + promoted_context_request_ids: frozenset[int] + ) -> Tuple[bool, Optional[SampleType]]: + """Agree on graph eligibility, rows and sampling tier across ADP ranks.""" + if not (self.enabled and self.config.enable_attention_dp + and self.config.mapping.tp_size > 1): + return True, None + local_sample_type = self._resolve_sample_type( + batch, promoted_context_request_ids) + graph_batch_info = self._gather_adp_graph_batch_info( + batch, can_run_cuda_graph, local_sample_type) + all_can_graph = all(rank_info[0] for rank_info in graph_batch_info) + all_batch_sizes_equal = all(rank_info[1] == graph_batch_info[0][1] + for rank_info in graph_batch_info) + if not all_can_graph or not all_batch_sizes_equal: + return False, None + forced_sample_type = (local_sample_type if all( + rank_info[2] == graph_batch_info[0][2] + for rank_info in graph_batch_info) else SampleType.FULL) + return True, forced_sample_type + @staticmethod def _get_mrope_position_delta(request: Any) -> Optional[Any]: mrope_position_delta = getattr(request, "py_mrope_position_delta", None) @@ -568,28 +668,10 @@ def maybe_get_cuda_graph( is_mixed_encoder_decoder = self._is_mixed_encoder_decoder_batch(batch) can_run_cuda_graph = self._can_run_cuda_graph_batch(batch) batch_size = batch.batch_size - # The sampling tier joins the graph key, so it has to agree across the - # attention-DP ranks or they would replay different graphs -- the same - # reason can_run_cuda_graph and batch_size are all-gathered here. Ranks - # see different batches (an idle one resolves FULL), so fall back to - # FULL for the whole group unless every rank picked the same tier. - forced_sample_type: Optional[SampleType] = None - if self.enabled and self.config.enable_attention_dp and self.config.mapping.tp_size > 1: - local_sample_type = self._resolve_sample_type( - batch, promoted_context_request_ids) - graph_batch_info = self._gather_adp_graph_batch_info( - batch, can_run_cuda_graph, local_sample_type) - all_can_run_cuda_graph = all(rank_info[0] - for rank_info in graph_batch_info) - all_batch_sizes_equal = all(rank_info[1] == graph_batch_info[0][1] - for rank_info in graph_batch_info) - - if not all_can_run_cuda_graph or not all_batch_sizes_equal: - return None, None, None - - forced_sample_type = (local_sample_type if all( - rank_info[2] == graph_batch_info[0][2] - for rank_info in graph_batch_info) else SampleType.FULL) + shapes_agree, forced_sample_type = self._agree_graph_shape( + batch, can_run_cuda_graph, promoted_context_request_ids) + if not shapes_agree: + return None, None, None if not self.enabled or not can_run_cuda_graph: return None, None, None @@ -989,6 +1071,25 @@ def _check_strict_buffer_stability(self, stored_meta: Dict[str, self._assert_graph_tensor_ptrs_stable( stored_meta["spec_metadata"], stored_meta["spec_metadata_ptrs"]) + def will_pad_to(self, padded_bs: int, num_requests: int) -> bool: + """Whether a batch of ``num_requests`` can actually reach ``padded_bs``. + + Mirrors the SIZE guards in :meth:`_get_padded_batch` only; the + resource-dependent bails (dummy KV capacity, encoder-decoder state) + are deliberately not predicted. The ragged fit sizes its bucket grid + from the padded row count, so a wrong yes strands the fitted total in + no captured bucket. + """ + if padded_bs <= 0 or num_requests <= 0: + return False + if padded_bs not in self.supported_batch_sizes: + return False + if padded_bs == num_requests: + return True # already at a captured size, nothing to add + # _get_padded_batch bails when padding_size + batch_size exceeds the + # configured batch size; that sum is just padded_bs. + return padded_bs <= self.config.batch_size + def _get_padded_batch(self, batch: ScheduledRequests, resource_manager: ResourceManager, runtime_draft_len: int) -> int: @@ -1015,20 +1116,7 @@ def _get_padded_batch(self, batch: ScheduledRequests, """ can_run_cuda_graph = self._can_run_cuda_graph_batch(batch) batch_size = batch.batch_size - new_batch_size = batch_size - - if self.enabled and self.config.enable_attention_dp and self.config.mapping.tp_size > 1: - # Padding only agrees on eligibility and batch size; the tier slot - # is unused here. It is still sent so both per-iteration exchanges - # share one fixed-size payload, and a constant keeps this path from - # resolving a tier for a batch it is about to change. - graph_batch_info = self._gather_adp_graph_batch_info( - batch, can_run_cuda_graph, SampleType.FULL) - all_can_run_cuda_graph = all(rank_info[0] - for rank_info in graph_batch_info) - if all_can_run_cuda_graph: - new_batch_size = max(rank_info[1] - for rank_info in graph_batch_info) + new_batch_size = self._get_padding_batch_size(batch, can_run_cuda_graph) if (not self.enabled or not self.padding_enabled or not can_run_cuda_graph @@ -1041,6 +1129,7 @@ def _get_padded_batch(self, batch: ScheduledRequests, # otherwise this reduces to plain batch-size rounding. padded_batch_size = self._round_up_batch_size_with_draft_len( new_batch_size, runtime_draft_len) + self._validate_padding_batch_size(batch, padded_batch_size) if batch_size == padded_batch_size: return 0 @@ -1065,12 +1154,56 @@ def _get_padded_batch(self, batch: ScheduledRequests, key=f"cuda_graph_padding_dummy_fallback_{runtime_draft_len}") return 0 + return self._append_padding_requests(batch, resource_manager, + runtime_draft_len, + padding_dummy_request, + padding_size) + + def _get_padding_batch_size(self, batch: ScheduledRequests, + can_run_cuda_graph: bool) -> int: + new_batch_size = batch.batch_size + if (self.enabled and self.config.enable_attention_dp + and self.config.mapping.tp_size > 1): + graph_batch_info = self._gather_adp_graph_batch_info( + batch, can_run_cuda_graph, SampleType.FULL) + if all(rank_info[0] for rank_info in graph_batch_info): + new_batch_size = max(rank_info[1] + for rank_info in graph_batch_info) + return new_batch_size + + def _validate_padding_batch_size(self, batch: ScheduledRequests, + padded_batch_size: int) -> None: + """Validate a specialized runner's pre-agreed padding shape.""" + + def _append_padding_requests(self, batch: ScheduledRequests, + resource_manager: ResourceManager, + runtime_draft_len: int, + padding_dummy_request: LlmRequest, + padding_size: int) -> int: batch.generation_requests.extend([padding_dummy_request] * padding_size) return padding_size - def _get_or_create_padding_dummy( - self, resource_manager: ResourceManager, - runtime_draft_len: int) -> Optional[LlmRequest]: + def _padding_dummy_cache(self, variant: int) -> Dict[int, LlmRequest]: + if variant != 0: + raise ValueError( + f"unsupported CUDA padding dummy variant {variant}") + return self.padding_dummy_requests + + def _padding_dummy_caches(self) -> Tuple[Dict[int, LlmRequest], ...]: + return (self.padding_dummy_requests, ) + + def reset_padding_dummies(self) -> None: + """Drop cached request objects when their resource manager changes.""" + self.padding_dummy_requests = {} + + def _on_padding_dummy_release(self) -> None: + """Invalidate specialized padding agreements after resource release.""" + + def _get_or_create_padding_dummy(self, + resource_manager: ResourceManager, + runtime_draft_len: int, + *, + variant: int = 0) -> Optional[LlmRequest]: """Returns the padding dummy request for the given draft length, allocating it from the KV cache manager on first use. @@ -1078,8 +1211,9 @@ def _get_or_create_padding_dummy( (e.g. saturated cache); the batch cannot be padded until a retry succeeds. """ - if runtime_draft_len in self.padding_dummy_requests: - return self.padding_dummy_requests[runtime_draft_len] + cache = self._padding_dummy_cache(variant) + if runtime_draft_len in cache: + return cache[runtime_draft_len] kv_cache_manager = resource_manager.get_resource_manager( self.config.kv_cache_manager_key) @@ -1104,8 +1238,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 +1258,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 +1331,27 @@ def release_padding_dummy(self, resource_manager: ResourceManager, manager that allocated part of it, and drops it from the runner so a later padded step re-creates it. - One dummy request ID is spread across up to four managers -- the main - KV cache manager, the one-model draft KV cache manager, the + Each dummy request ID is spread across up to four managers -- + the main KV cache manager, the one-model draft KV cache manager, the speculative resource manager slot and the encoder-decoder cross-KV - cache manager. Releasing only the main one leaves the others holding - the ID, and re-creation reuses the same - ``CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len``. + cache manager. Releasing only the main one leaves the others holding + the ID; re-creation deterministically reuses the same variant ID. Returns True if a dummy was held for that draft length. """ - dummy_request = self.padding_dummy_requests.pop(runtime_draft_len, None) - if dummy_request is None: + dummy_requests = [ + cache.pop(runtime_draft_len, None) + for cache in self._padding_dummy_caches() + ] + dummy_requests = [ + request for request in dummy_requests if request is not None + ] + if not dummy_requests: return False - for manager in self._padding_dummy_managers(resource_manager): - manager.free_resources(dummy_request) + self._on_padding_dummy_release() + for dummy_request in dummy_requests: + for manager in self._padding_dummy_managers(resource_manager): + manager.free_resources(dummy_request) return True def _can_pad_any_batch(self, runtime_draft_len: int) -> bool: @@ -1241,6 +1410,12 @@ def preallocate_padding_dummies(self, logger.info( "Pre-allocated the CUDA graph padding dummy request " f"(draft_len={draft_len}) at warmup.") + self._preallocate_extra_padding_dummy(resource_manager, draft_len) + + def _preallocate_extra_padding_dummy(self, + resource_manager: ResourceManager, + draft_len: int) -> None: + """Reserve any additional padding variants required by a subclass.""" def _add_cross_dummy_request( self, dummy_request: LlmRequest, resource_manager: ResourceManager, @@ -1264,11 +1439,6 @@ def _add_cross_dummy_request( if cross_dummy_requests is not None: return True - kv_cache_manager = resource_manager.get_resource_manager( - self.config.kv_cache_manager_key) - kv_cache_manager.free_resources(dummy_request) - if draft_kv_cache_manager is not None: - draft_kv_cache_manager.free_resources(dummy_request) return False @staticmethod @@ -1338,12 +1508,207 @@ def clear(self): self.graph_outputs.clear() self.graph_metadata.clear() self._pre_replay_position_ids = None - self.padding_dummy_requests = {} + self.reset_padding_dummies() del self.memory_pool self.memory_pool = None torch.cuda.empty_cache() +class DSparkCUDAGraphRunner(CUDAGraphRunner): + """Ragged verification graph policy for confidence-enabled DSpark.""" + + def __init__(self, config: CUDAGraphRunnerConfig) -> None: + super().__init__(config) + # Zero-real exact cells use low/high rows to represent V=G*q+r. + self.ragged_pad_verify_len = 0 + self.ragged_zero_real_high_rows = 0 + # The decoder publishes the fitted bucket before graph lookup. + self.agreed_ragged_bucket: Optional[int] = None + self.adp_shape_agreement: Optional[ADPShapeAgreement] = None + self.secondary_padding_dummy_requests: Dict[int, LlmRequest] = {} + + def _ragged_verify_bucket(self, batch: ScheduledRequests) -> Optional[int]: + """Return the agreed token bucket only when every row has a window.""" + # A window-less batch must not replay stale token-major row maps. + if any( + getattr(request, "py_verify_len", None) is None + for request in batch.generation_requests): + return None + # Use the globally fitted bucket, never a rank-local window sum. + return self.agreed_ragged_bucket + + @staticmethod + def _local_draft_len(batch: ScheduledRequests) -> int: + """Report malformed batches' maximum so the shape gate declines them.""" + return max((len(request.py_draft_tokens) + for request in batch.generation_requests), + default=0) + + def _agree_graph_shape( + self, batch: ScheduledRequests, can_run_cuda_graph: bool, + promoted_context_request_ids: frozenset[int] + ) -> Tuple[bool, Optional[SampleType]]: + agreement = self.adp_shape_agreement + # The policy agreement covers shape, not the in-graph sampling tier. + reuse_agreement = bool( + agreement is not None and agreement.reuse_graph_shape + and not self.enable_in_graph_sampling + and agreement.matches(batch, agreement.iteration, padded=True)) + if reuse_agreement: + if (not can_run_cuda_graph + or batch.batch_size != agreement.graph_batch_size or + self._ragged_verify_bucket(batch) != agreement.ragged_bucket + or self._local_draft_len(batch) != agreement.draft_len): + raise RuntimeError( + "DSpark cached ADP graph-shape agreement no longer " + "matches the local replay key") + adp_graph_agreement = (self.enabled and self.config.enable_attention_dp + and self.config.mapping.tp_size > 1) + local_sample_type = self._resolve_sample_type( + batch, + promoted_context_request_ids) if adp_graph_agreement else None + if not adp_graph_agreement or reuse_agreement: + return True, None + + # Replay freezes all_rank_num_tokens. Compare every rank-local shape + # component before launch; -1 denotes a non-applicable field. + local_bucket = self._ragged_verify_bucket( + batch) if can_run_cuda_graph else None + peer_shapes = self.config.dist.tp_allgather_int64([ + can_run_cuda_graph, + batch.batch_size, + -1 if local_bucket is None else int(local_bucket), + self._local_draft_len(batch) if can_run_cuda_graph else -1, + local_sample_type.value, + ]).tolist() + if (not all(peer[0] for peer in peer_shapes) + or not all(peer[1:4] == peer_shapes[0][1:4] + for peer in peer_shapes)): + return False, None + forced_sample_type = (local_sample_type if all( + peer[4] == peer_shapes[0][4] + for peer in peer_shapes) else SampleType.FULL) + return True, forced_sample_type + + def _get_num_tokens_for_key(self, key: KeyType) -> int: + if key.ragged_verify_bucket is not None: + if key.num_contexts: + raise ValueError( + "Ragged verifier graphs do not support mixed context batches." + ) + return int(key.ragged_verify_bucket) + return super()._get_num_tokens_for_key(key) + + def _reusable_padding_agreement( + self, batch: ScheduledRequests) -> Optional[ADPShapeAgreement]: + agreement = self.adp_shape_agreement + if (agreement is not None and agreement.reuse_graph_shape and + agreement.matches(batch, agreement.iteration, padded=False)): + return agreement + return None + + def _get_padding_batch_size(self, batch: ScheduledRequests, + can_run_cuda_graph: bool) -> int: + agreement = self._reusable_padding_agreement(batch) + if agreement is None: + return super()._get_padding_batch_size(batch, can_run_cuda_graph) + if not can_run_cuda_graph or not agreement.padding_ready: + raise RuntimeError( + "DSpark cached ADP padding agreement lost a local " + "eligibility invariant") + return agreement.widest_batch_size + + def _validate_padding_batch_size(self, batch: ScheduledRequests, + padded_batch_size: int) -> None: + agreement = self._reusable_padding_agreement(batch) + if (agreement is not None + and padded_batch_size != agreement.graph_batch_size): + raise RuntimeError( + "DSpark cached ADP padding agreement selected a different " + "graph batch size") + + def _append_padding_requests(self, batch: ScheduledRequests, + resource_manager: ResourceManager, + runtime_draft_len: int, + padding_dummy_request: LlmRequest, + padding_size: int) -> int: + # The ragged fit chooses windows before cached dummies are appended. + secondary_padding_dummy = None + high_padding_rows = 0 + if self.ragged_zero_real_high_rows: + if (batch.batch_size != 1 + or not batch.generation_requests[0].is_attention_dp_dummy + or self.ragged_zero_real_high_rows + > batch.batch_size + padding_size): + raise RuntimeError( + "DSpark zero-real heterogeneous padding lost its single " + "scheduled attention-DP dummy invariant") + high_padding_rows = self.ragged_zero_real_high_rows - 1 + if high_padding_rows > 0: + secondary_padding_dummy = self._get_or_create_padding_dummy( + resource_manager, runtime_draft_len, variant=1) + if secondary_padding_dummy is None: + logger.warning_once( + "Failed to allocate the secondary DSpark CUDA graph " + "padding dummy; falling back to native verification.", + key=("dspark_secondary_padding_dummy_fallback_" + f"{runtime_draft_len}")) + return 0 + ragged = any( + getattr(request, "py_verify_len", None) is not None + for request in batch.generation_requests) + padding_dummy_request.py_verify_len = (self.ragged_pad_verify_len + if ragged else None) + if secondary_padding_dummy is not None: + secondary_padding_dummy.py_verify_len = ( + self.ragged_pad_verify_len + 1 if ragged else None) + low_padding_rows = padding_size - high_padding_rows + if low_padding_rows < 0: + raise RuntimeError( + "DSpark zero-real high-row count exceeds CUDA padding rows") + if high_padding_rows: + batch.generation_requests.extend([secondary_padding_dummy] * + high_padding_rows) + batch.generation_requests.extend([padding_dummy_request] * + low_padding_rows) + return padding_size + + def _padding_dummy_cache(self, variant: int) -> Dict[int, LlmRequest]: + if variant == 1: + return self.secondary_padding_dummy_requests + return super()._padding_dummy_cache(variant) + + def _padding_dummy_caches(self) -> Tuple[Dict[int, LlmRequest], ...]: + return super()._padding_dummy_caches() + ( + self.secondary_padding_dummy_requests, ) + + def reset_padding_dummies(self) -> None: + super().reset_padding_dummies() + self.secondary_padding_dummy_requests = {} + self.ragged_zero_real_high_rows = 0 + self.adp_shape_agreement = None + + def _on_padding_dummy_release(self) -> None: + # Releasing a dummy invalidates the agreement that certified it. + self.adp_shape_agreement = None + + def _preallocate_extra_padding_dummy(self, + resource_manager: ResourceManager, + draft_len: int) -> None: + if not getattr(self.spec_config, "confidence_sps_table_path", None): + return + if self._get_or_create_padding_dummy(resource_manager, + draft_len, + variant=1) is None: + logger.warning("Could not pre-allocate the secondary DSpark CUDA " + "graph padding dummy " + f"(draft_len={draft_len}); non-divisible zero-real " + "exact cells will remain native.") + else: + logger.info("Pre-allocated the secondary DSpark CUDA graph " + f"padding dummy (draft_len={draft_len}).") + + EncoderKeyType: TypeAlias = Tuple[int, int, int] _ENCODER_SOURCE_SEQ_LENS = "_encoder_source_seq_lens" _ENCODER_SOURCE_TO_SLOT = "_encoder_source_to_slot" diff --git a/tensorrt_llm/_torch/pyexecutor/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/dspark_observation_runtime.py b/tensorrt_llm/_torch/pyexecutor/dspark_observation_runtime.py new file mode 100644 index 000000000000..d942ff8cf789 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/dspark_observation_runtime.py @@ -0,0 +1,341 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Explicit offline observation factory over normal native/automatic execution. + +No fixed-cell/profile-only engine is used here. Calibration alone evaluates the +confidence head under native full-K verification and copies observations before +the SAME normal sampler event; benchmark arms add no calibration copies. +""" + +from __future__ import annotations + +from functools import partial +from typing import Any + +from tensorrt_llm.tools.dspark_observations import NativePairs, ProposalSnapshot + + +class ObservationExecutorMixin: + """Thin lifecycle joins; every forward, sampler/update and error delegates.""" + + def __init__(self, *args, **kwargs): + self._observation_inputs = {} + self._observation_states = {} + self._observation_pending = {} + self._observation_forward = None + self._observation_failures = [] + super().__init__(*args, **kwargs) + self._native_pairs = NativePairs(self.model_engine.spec_config.max_draft_len) + + def _handle_errors( + self, + error_msg=None, + *, + requests=None, + charge_budget=True, + fatal_is_collective_aligned=False, + ): + self._observation_failures.append(str(error_msg or "Original executor error")) + return super()._handle_errors( + error_msg, + requests=requests, + charge_budget=charge_budget, + fatal_is_collective_aligned=fatal_is_collective_aligned, + ) + + def _forward_step(self, batch, new_tensors_device=None): + if self.is_warmup or not self.model_engine._collect_calibration: + return super()._forward_step(batch, new_tensors_device) + if self._observation_forward is not None: + raise RuntimeError("Overlapping unpaired observation forwards") + requests = tuple(r for r in batch.generation_requests if not r.is_dummy) + source = self._observation_inputs.get(id(new_tensors_device)) + input_host = None + if source is not None: + if source.state.device is not new_tensors_device: + raise RuntimeError("Original device input incarnation changed") + # Issued before original forward can reuse persistent output storage. + # The next original sample_async event covers this same-stream copy. + input_host = self.sampler._copy_to_host(new_tensors_device.next_draft_tokens) + elif requests: + self._native_pairs.stats["bootstrap_or_no_bound_input"] += len(requests) + # Existing source stamp scalar is observation identity only. It does + # not alter proposal/acceptance math or construct a scheduling policy. + draft_sequence = self.model_engine.model.spec_worker.bump_draft_seq() + 1 + self._observation_forward = ( + self.iter_counter, + requests, + source, + input_host, + draft_sequence, + ) + result = super()._forward_step(batch, new_tensors_device) + if result is None: + raise RuntimeError("Original forward failed; observations cannot be promoted") + return result + + def _sample_async(self, batch, output): + if self.is_warmup or not self.model_engine._collect_calibration: + return super()._sample_async(batch, output) + worker = self.model_engine.model.spec_worker + confidence = worker.staged_confidence_buffer() + if confidence is None or self._observation_forward is None: + raise RuntimeError("Offline full-K collector has no produced confidence/forward") + iteration, requests, source, input_host, draft_sequence = self._observation_forward + rows = {} + for request in requests: + row = worker.confidence_row_for(request.py_request_id) + if row is None or request.py_seq_slot is None: + raise RuntimeError("Original request has no worker confidence/sampler slot") + rows[request.py_request_id] = (request, int(request.py_seq_slot), int(row)) + # Already-produced full-K logits are copied with the original sampler's + # helper BEFORE it records its normal event. No extra event or wait. + host = self.sampler._copy_to_host(confidence) + stamps = self.sampler._copy_to_host(worker.confidence_stamp_buffer()) + state = super()._sample_async(batch, output) + if state is None or id(state) in self._observation_pending: + raise RuntimeError("Missing or reused original sampler state") + snapshot = ProposalSnapshot(iteration, state, host, rows, stamps, draft_sequence) + self._observation_inputs[id(state.device)] = snapshot + self._observation_states[id(state)] = snapshot + self._observation_pending[id(state)] = (iteration, requests, source, input_host) + self._observation_forward = None + return state + + def _update_requests(self, state, resource_manager=None): + if self.is_warmup or not self.model_engine._collect_calibration: + return super()._update_requests(state, resource_manager) + pending = self._observation_pending.get(id(state)) + snapshot = self._observation_states.get(id(state)) + if pending is None or snapshot is None or snapshot.state is not state: + raise RuntimeError("Unknown or consumed observation sampler state") + before = {id(request): request.py_decoding_iter for request in state.requests} + result = super()._update_requests(state, resource_manager) + # Original SpecSampler.update_requests performed its existing synchronize + # and updates. Parent errors are fatal to publication, not waived here. + if self._observation_failures: + raise RuntimeError("Original sampler/update error; no paired-data publication") + snapshot.updated = True + iteration, requests, source, input_host = pending + by_object = {id(request): index for index, request in enumerate(state.requests)} + for request in requests: + index = by_object.get(id(request)) + if index is None: + raise RuntimeError("Original generation request absent from sampler state") + if request.py_decoding_iter <= before[id(request)]: + self._native_pairs.stats["retired_or_not_advanced"] += 1 + continue + if source is None: + continue + if request.py_request_id not in source.requests: + # A finished-context output has only placeholder draft rows, + # not a confidence-scored block from a generation forward. + self._native_pairs.stats["prefill_placeholder_no_scored_proposal"] += 1 + continue + slot = int(request.py_seq_slot) + token_window = ( + int(state.host.verify_lens[index]) + if state.host.verify_lens is not None + else state.runtime_draft_len + 1 + ) + verified = min(int(state.draft_lens[index]), token_window - 1) + if verified != self._native_pairs.physical_k: + # First-draft bootstrap has no scored full-K proposal. Never + # reinterpret its placeholder confidence as calibration data. + self._native_pairs.stats["bootstrap_or_partial"] += 1 + continue + self._native_pairs.record( + source=source, + request=request, + target_iteration=iteration, + input_proposal=input_host[slot], + target_slot=slot, + accepted=int(state.host.new_tokens_lens[slot]) - 1, + verified=verified, + ) + del self._observation_pending[id(state)] + return result + + def _observation_drain_predicates(self): + previous = self.previous_batch + state = previous.sample_state if previous is not None else None + snapshot = self._observation_states.get(id(state)) if state is not None else None + return { + "failures": len(self._observation_failures), + "active_requests": len(self.active_requests), + "waiting_requests": len(self.waiting_queue), + "pending_samples": len(self._observation_pending), + "forward_pending": self._observation_forward is not None, + "previous_batch_present": previous is not None, + "pending_matches_previous": ( + state is not None + and len(self._observation_pending) == 1 + and id(state) in self._observation_pending + and snapshot is not None + and snapshot.state is state + and not snapshot.updated + ), + } + + def observation_receipt(self, *, reset=False, control=None): + with self.control_action(drain=True): + before = self._observation_drain_predicates() + consumed = False + if self.model_engine._collect_calibration: + if control is None: + raise RuntimeError("Offline drain needs the original gang-control communicator") + peers = control.allgather(before) + quiescent = all( + row["failures"] == 0 + and row["active_requests"] == 0 + and row["waiting_requests"] == 0 + and not row["forward_pending"] + for row in peers + ) + no_tail = all(row["pending_samples"] == 0 for row in peers) + exact_tail = all( + row["pending_samples"] == 1 and row["pending_matches_previous"] for row in peers + ) + if not quiescent or not (no_tail or exact_tail): + raise RuntimeError( + "Original requests/updates did not naturally drain: " + f"gang_predicates={peers!r}" + ) + if exact_tail: + # Every rank is paused at the original exclusive boundary. + # Complete the same existing sampler/update/resource tail + # before checking the observer's stronger publication gate. + self._consume_previous_batch_for_rebalance() + consumed = True + after = self._observation_drain_predicates() + if ( + self._observation_failures + or self.active_requests + or len(self.waiting_queue) + or self._observation_pending + or self._observation_forward is not None + ): + raise RuntimeError( + "Original requests/updates did not naturally drain: " + f"before={before!r}, after={after!r}" + ) + runner = self.model_engine.cuda_graph_runner + result = { + "drain": { + "before": before, + "after": after, + "original_previous_batch_consumed": consumed, + }, + "routing": dict(runner.observed_routes), + "graph_keys": list(runner.observed_keys.values()), + "resolved_confidence_enabled": bool( + self.model_engine.spec_config.enable_confidence_scheduling + ), + "runtime_confidence_enabled": bool( + getattr(self.model_engine, "_dspark_confidence_enabled", False) + ), + } + if reset: + runner.observed_routes = {"native_replays": 0, "compact_replays": 0} + runner.observed_keys.clear() + return result + + +def create_observation_executor( + *, llm_args, collect_calibration=False, checkpoint_dir=None, tokenizer=None +): + """Only this public offline caller gets observation hooks; normal serving is unchanged.""" + from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( + CUDAGraphRunner, + DSparkCUDAGraphRunner, + KeyType, + ) + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner + from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + from tensorrt_llm._torch.pyexecutor.py_executor_creator import ( + _RuntimeConstructors, + create_py_executor, + ) + from tensorrt_llm._torch.speculative.dspark import DSv4DSparkWorker, resolve_confidence_head + from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig + + spec = llm_args.speculative_config + if not isinstance(spec, DSparkDecodingConfig) or not spec.draft_is_embedded_in_target: + raise ValueError("The DSpark adapter requires embedded chain drafting") + if collect_calibration and (spec.enable_confidence_scheduling or spec.draft_len_schedule): + raise ValueError( + "Offline calibration requires native full-K, not policy-trimmed verification" + ) + + # Retain PR5's ordinary native-OFF versus DSpark-ON graph implementation. + graph_base = DSparkCUDAGraphRunner if spec.enable_confidence_scheduling else CUDAGraphRunner + + class ObservedGraphs(graph_base): + def __init__(self, *args, **kwargs): + self.observed_routes = {"native_replays": 0, "compact_replays": 0} + self.observed_keys = {} + super().__init__(*args, **kwargs) + + def replay(self, key, inputs, *, position_id_offsets=None): + result = super().replay(key, inputs, position_id_offsets=position_id_offsets) + key = KeyType(*key) # same accepted tuple/KeyType contract as the original runner + if key.is_first_draft or key.draft_len != spec.max_draft_len: + return result + raw = dict(key._asdict()) + # Existing KeyType carries a SampleType enum; no arbitrary objects + # or environment values are recorded by this finite graph witness. + raw["sample_type"] = str(raw["sample_type"]) + budget = getattr(key, "ragged_verify_bucket", None) + compact = budget is not None and 0 < budget < key.batch_size * (spec.max_draft_len + 1) + self.observed_routes["compact_replays" if compact else "native_replays"] += 1 + self.observed_keys[repr(key)] = raw + return result + + class ObservationDecoderRunner(DecoderRunner): + def _initialize_cuda_graph_runner(self): + if collect_calibration: + worker = self.model.spec_worker + if type(worker) is not DSv4DSparkWorker or worker._win_inited: + raise RuntimeError( + "Calibrated head must be enabled before original graph/storage setup" + ) + if resolve_confidence_head(self.model.draft_model) is None: + raise ValueError("Actual drafter has no confidence head") + worker.return_confidence = True # head evaluation only, no planner/selection + return super()._initialize_cuda_graph_runner() + + def _create_cuda_graph_runner(self, config): + return ObservedGraphs(config) + + class ObservationEngine(PyTorchModelEngine): + def __init__(self, **kwargs: Any): + self._collect_calibration = collect_calibration + super().__init__(**kwargs) + + def _initialize_decoder_runner(self, runner_cls, runner_config=None, **runner_kwargs): + if runner_cls is not DecoderRunner: + raise ValueError("Offline DSpark observation requires the original decoder runner") + return super()._initialize_decoder_runner( + ObservationDecoderRunner, runner_config, **runner_kwargs + ) + + class ObservationExecutor(ObservationExecutorMixin, PyExecutor): + pass + + constructors = _RuntimeConstructors( + model_engine=ObservationEngine, + executor=ObservationExecutor, + kv_cache_creator=KvCacheCreator, + ) + return create_py_executor( + llm_args=llm_args, + checkpoint_dir=checkpoint_dir, + tokenizer=tokenizer, + runtime_constructors=constructors, + ) + + +def observation_factory(*, collect_calibration=False): + return partial(create_observation_executor, collect_calibration=collect_calibration) diff --git a/tensorrt_llm/_torch/pyexecutor/dspark_profile_engine.py b/tensorrt_llm/_torch/pyexecutor/dspark_profile_engine.py new file mode 100644 index 000000000000..ecf50df546cf --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/dspark_profile_engine.py @@ -0,0 +1,462 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Non-serving fixed-cell DSpark profiling on the original model engine.""" + +from __future__ import annotations + +import hashlib +import os +from collections.abc import Sequence +from typing import Any + +import torch + +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( + CUDAGraphRunner, + CUDAGraphRunnerConfig, + DSparkCUDAGraphRunner, +) +from tensorrt_llm._torch.pyexecutor.dspark_profile_graph_storage import ( + snapshot_metadata_storage, + validate_metadata_storage, +) +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner +from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager +from tensorrt_llm._torch.pyexecutor.sampler import SampleStateTensors +from tensorrt_llm._torch.pyexecutor.sampler.sampler_common import SampleType +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.speculative.dspark import DSv4DSparkWorker, resolve_confidence_head +from tensorrt_llm._torch.speculative.dspark_cost_profile import ( + DSparkCostProfileConfig, + DSparkProfileCell, +) +from tensorrt_llm._torch.speculative.spec_sampler_base import SampleStateTensorsSpec +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig, PrefillCudaGraphBackend +from tensorrt_llm.tools.verify_costs import EventLeasePool + +__all__ = ["DSparkCostProfileEngine"] + + +def _require_shape_agreement(rows: Sequence[Sequence[int]]) -> None: + if not rows or any(len(row) != 5 or row[0] != 1 for row in rows): + raise ValueError("Every profile rank must provide a valid decode-only cell") + if any(tuple(row[1:]) != tuple(rows[0][1:]) for row in rows): + raise ValueError("Profile ranks disagree on G, physical K, V or replay ordinal") + + +def _profile_contract_row(profile: DSparkCostProfileConfig, valid_ladder: bool) -> list[int]: + digest = hashlib.sha256(profile.model_dump_json().encode()).digest() + return [ + int(valid_ladder), + profile.physical_k, + int(profile.record_confidence_head), + profile.max_samples, + *(int.from_bytes(digest[index : index + 8], signed=True) for index in range(0, 32, 8)), + ] + + +def _require_profile_agreement(rows: Sequence[Sequence[int]]) -> None: + if ( + not rows + or any(len(row) != 8 or row[0] != 1 for row in rows) + or any(tuple(row[1:]) != tuple(rows[0][1:]) for row in rows) + ): + raise ValueError( + "Every profile rank must agree on the construction-only grid and head mode" + ) + + +def _valid_previous_state(state: Any, requests: Sequence[Any], physical_k: int) -> bool: + """Validate original overlap tensor extents without reading device values.""" + if state is None: + return all(len(request.py_draft_tokens) == physical_k for request in requests) + if not isinstance(state, SampleStateTensorsSpec): + return False + tensors = (state.new_tokens, state.new_tokens_lens, state.next_draft_tokens) + if any( + not isinstance(value, torch.Tensor) + or not value.is_cuda + or value.dtype not in (torch.int32, torch.int64) + or not value.is_contiguous() + for value in tensors + ): + return False + tokens, lengths, drafts = tensors + if ( + tokens.ndim != 3 + or tokens.shape[0] != physical_k + 1 + or tokens.shape[2] != 1 + or lengths.ndim != 1 + or drafts.ndim != 2 + or drafts.shape[1] != physical_k + or tokens.shape[1] != lengths.shape[0] + or drafts.shape[0] != lengths.shape[0] + or len({value.device for value in tensors}) != 1 + ): + return False + slots = [getattr(request, "py_batch_idx", None) for request in requests] + return all(type(slot) is int and 0 <= slot < lengths.shape[0] for slot in slots) and len( + set(slots) + ) == len(slots) + + +class _DSparkProfileWorker(DSv4DSparkWorker): + def __init__( + self, + spec_config: DSparkDecodingConfig, + mapping: Any, + *, + record_head: bool, + use_separate_draft_kv_cache: bool, + ) -> None: + super().__init__(spec_config, mapping, use_separate_draft_kv_cache) + self.return_confidence = record_head + + def _build_verify_planner(self, draft_model: torch.nn.Module, block_size: int) -> None: + if block_size != self.spec_config.max_draft_len: + raise ValueError("Profile head width must match the physical draft block") + if resolve_confidence_head(draft_model) is None: + raise ValueError("Head-ON profiling requires the loaded confidence head") + # Buffers and the original head forward are retained; no serving planner is constructed. + self.verify_planner = None + + +class _DSparkProfileGraphs(DSparkCUDAGraphRunner): + _snapshot_graph_tensor_ptrs = staticmethod(snapshot_metadata_storage) + _assert_graph_tensor_ptrs_stable = staticmethod(validate_metadata_storage) + + def __init__(self, config: CUDAGraphRunnerConfig, profile: DSparkCostProfileConfig) -> None: + super().__init__(config) + self._profile = profile + self._dspark_trims_submitted_tokens = True + self.active_profile_cell: DSparkProfileCell | None = None + self.active_request_ids: tuple[int, ...] = () + self._profile_pool = EventLeasePool(lambda: torch.cuda.Event(enable_timing=True)) + self._profile_pool.initialize() + torch.cuda.current_stream().synchronize() + self._profile_records: list[dict[str, Any]] = [] + self._profile_ptrs: dict[Any, tuple[dict[str, Any], dict[str, Any]]] = {} + + def capture(self, key: Any, *args: Any, **kwargs: Any) -> Any: + output = super().capture(key, *args, **kwargs) + metadata = self.graph_metadata[key] + self._profile_ptrs[key] = ( + self._snapshot_graph_tensor_ptrs(metadata["attn_metadata"]), + self._snapshot_graph_tensor_ptrs(metadata["spec_metadata"]) + if metadata["spec_metadata"] is not None + else {}, + ) + return output + + def maybe_get_cuda_graph(self, batch: ScheduledRequests, *args: Any, **kwargs: Any) -> Any: + result = super().maybe_get_cuda_graph(batch, *args, **kwargs) + cell = self.active_profile_cell + if cell is not None: + key = result[2] + if key is None or key not in self.graphs: + raise RuntimeError("A real profile cell requires its previously captured graph") + expected_v = None if cell.verifier_budget == 0 else cell.verifier_budget + if (key.batch_size, key.draft_len, key.ragged_verify_bucket) != ( + cell.graph_batch_size, + self._profile.physical_k, + expected_v, + ): + raise RuntimeError("Observed replay key does not match the declared profile cell") + if key.num_contexts != 0 or not key.is_all_greedy_sample: + raise RuntimeError("The fixed-cell pilot measures greedy decode-only graphs") + return result + + def replay( + self, + key: Any, + current_inputs: dict[str, Any], + *, + position_id_offsets: torch.Tensor | None = None, + ) -> Any: + cell = self.active_profile_cell + if cell is None: + return super().replay(key, current_inputs, position_id_offsets=position_id_offsets) + if len(self._profile_records) >= self._profile.max_samples: + raise RuntimeError("Preallocated profile timing capacity exhausted") + metadata = self.graph_metadata[key] + attn_ptrs, spec_ptrs = self._profile_ptrs[key] + storage_context = { + "rank": self.config.mapping.rank, + "graph_key": tuple( + value.name if isinstance(value, SampleType) else value for value in key + ), + } + self._assert_graph_tensor_ptrs_stable( + metadata["attn_metadata"], attn_ptrs, context=storage_context + ) + if metadata["spec_metadata"] is not None: + self._assert_graph_tensor_ptrs_stable( + metadata["spec_metadata"], spec_ptrs, context=storage_context + ) + lease = self._profile_pool.begin() + output = super().replay(key, current_inputs, position_id_offsets=position_id_offsets) + self._profile_pool.end(lease) + self._profile_records.append( + { + "sample_index": len(self._profile_records), + "_event_lease": lease, + "rank": self.config.mapping.rank, + "graph_batch_size": cell.graph_batch_size, + "physical_k": self._profile.physical_k, + "verifier_budget": cell.verifier_budget, + "request_ids": self.active_request_ids, + "graph_key": tuple( + value.name if isinstance(value, SampleType) else value for value in key + ), + } + ) + return output + + def complete_sample(self, index: int) -> None: + """Release events after the original sampler/update completion boundary.""" + record = self._profile_records[index] + lease = record["_event_lease"] + elapsed = self._profile_pool.complete(lease) + del record["_event_lease"] + record["replay_entry_elapsed_ms"] = elapsed + + def completed_profile_samples(self) -> list[dict[str, Any]]: + """Return completed raw measurements without waiting or inventing costs.""" + if self._profile_pool.pending: + raise RuntimeError("Profile generation still owns unfinished event leases") + if any("replay_entry_elapsed_ms" not in row for row in self._profile_records): + raise RuntimeError("Profile replay has no original update completion") + return [dict(row) for row in self._profile_records] + + +class _DSparkProfileDecoderRunner(DecoderRunner): + """Offline-only capture/state owner, using the current decoder lifecycle.""" + + def __init__(self, model, config, *, profile, **runner_kwargs): + self._cost_profile = profile + super().__init__(model, config, **runner_kwargs) + + @property + def _ragged_verification_enabled(self) -> bool: + return True + + def _initialize_cuda_graph_runner(self): + old_worker = self.model.spec_worker + if type(old_worker) is not DSv4DSparkWorker or old_worker._win_inited: + raise RuntimeError("Profile worker must be installed before its first initialization") + if old_worker.staged_confidence_buffer() is not None or self.attn_metadata is not None: + raise RuntimeError( + "Profile construction cannot replace initialized state or graph addresses" + ) + matches = [index for index, entry in enumerate(self.model.epilogue) if entry is old_worker] + if len(matches) != 1: + raise RuntimeError("Original speculative worker must have one epilogue owner") + worker = _DSparkProfileWorker( + self._config.spec_config, + self.mapping, + record_head=self._cost_profile.record_confidence_head, + use_separate_draft_kv_cache=self.model.use_separate_draft_kv_cache, + ) + worker.set_draft_model(self.model.draft_model) + self.model.spec_worker = worker + self.model.epilogue[matches[0]] = worker + self._dspark_trims_submitted_tokens = True + self.previous_verify_lens = [] + return super()._initialize_cuda_graph_runner() + + def _create_cuda_graph_runner(self, config: CUDAGraphRunnerConfig) -> CUDAGraphRunner: + return _DSparkProfileGraphs(config, self._cost_profile) + + def _get_graphs_to_capture(self, cuda_graph_batch_sizes: list[int]) -> list[tuple[int, ...]]: + return self._cost_profile.graph_shapes(cuda_graph_batch_sizes) + + def ragged_verify_token_buckets(self, padded_bs: int) -> list[int]: + return sorted( + cell.verifier_budget + for cell in self._cost_profile.cells + if cell.graph_batch_size == padded_bs and cell.verifier_budget > 0 + ) + + +class DSparkCostProfileEngine(PyTorchModelEngine): + """Explicit offline engine; the ordinary public serving factories cannot create it. + + Args: + profile: Validated construction-only cell grid. + engine_kwargs: Original PyTorchModelEngine construction arguments. The + DSpark serving configuration must keep confidence scheduling off. + """ + + def __init__(self, *, profile: DSparkCostProfileConfig, **engine_kwargs: Any) -> None: + self._cost_profile = profile + self._profile_call_active = False + spec = engine_kwargs.get("spec_config") + if not isinstance(spec, DSparkDecodingConfig): + raise TypeError("Profile-only construction requires DSparkDecodingConfig") + if spec.max_draft_len != profile.physical_k or not spec.draft_is_embedded_in_target: + raise ValueError("Profile-only construction requires embedded DSpark with matching K") + if ( + spec.enable_confidence_scheduling + or spec.enable_fused_confidence_scheduler + or getattr(spec, "confidence_mixed_evidence", None) is not None + or spec.draft_len_schedule is not None + ): + raise ValueError("Profile-only construction cannot enable serving confidence policy") + if any( + getattr(spec, name, None) + for name in ( + "confidence_sps_table_path", + "confidence_sps_live_fingerprint_path", + "confidence_admission_receipt_path", + "confidence_admission_receipt_sha256", + "confidence_sts_path", + ) + ): + raise ValueError( + "Profile-only construction does not consume serving admission artifacts" + ) + if os.environ.get("TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS"): + raise ValueError("Fixed-cell costs require real acceptance, not forced acceptance") + super().__init__(**engine_kwargs) + + def _initialize_decoder_runner(self, runner_cls, runner_config=None, **runner_kwargs): + valid_ladder = True + try: + self._cost_profile.validate_graph_ladder(self._cuda_graph_batch_sizes) + except ValueError: + valid_ladder = False + if self.mapping.tp_size > 1 and self.dist is None: + raise ValueError("Multi-rank profiling requires the original distributed communicator") + row = _profile_contract_row(self._cost_profile, valid_ladder) + rows = self.dist.tp_allgather_int64(row).tolist() if self.mapping.tp_size > 1 else [row] + _require_profile_agreement(rows) + if ( + runner_cls is not DecoderRunner + or self.is_multimodal + or self._is_encoder_decoder_model() + or self.max_beam_width != 1 + or self.mapping.pp_size != 1 + or self.mapping.cp_size != 1 + or self.cuda_graph_config is None + or not self._cuda_graph_padding_enabled + or self.prefill_cuda_graph_backend != PrefillCudaGraphBackend.DISABLED + or self.llm_args.lora_config is not None + or self.llm_args.enable_in_graph_sampling + ): + raise ValueError("Unsupported profile-only model, topology or graph configuration") + return super()._initialize_decoder_runner( + _DSparkProfileDecoderRunner, + runner_config, + profile=self._cost_profile, + **runner_kwargs, + ) + + def _set_warmup_ragged_windows(self, batch, verify_bucket, draft_len): + if not isinstance(self._runner, _DSparkProfileDecoderRunner): + raise RuntimeError("Prescribed windows require the offline decoder owner") + return self._runner._set_warmup_ragged_windows(batch, verify_bucket, draft_len) + + def forward(self, *args: Any, **kwargs: Any) -> Any: + raise RuntimeError("Profile-only engines cannot execute the ordinary serving forward API") + + def forward_profile_prefill( + self, batch: ScheduledRequests, resource_manager: ResourceManager + ) -> Any: + """Initialize genuine request KV using the original unmeasured context forward.""" + valid = ( + batch.num_context_requests > 0 + and not batch.num_generation_requests + and not batch.num_encoder_requests + and not self._profile_call_active + ) + row = [int(valid), -1, self._cost_profile.physical_k, 0, -1] + rows = self.dist.tp_allgather_int64(row).tolist() if self.mapping.tp_size > 1 else [row] + _require_shape_agreement(rows) + return super().forward(batch, resource_manager=resource_manager) + + def forward_profile_cell( + self, + cell: DSparkProfileCell, + batch: ScheduledRequests, + resource_manager: ResourceManager, + *, + new_tensors_device: SampleStateTensors | None = None, + cache_indirection_buffer: torch.Tensor | None = None, + ) -> Any: + """Borrow real scheduled state and run one declared cell through the original forward. + + The caller retains ownership of request/sampler/KV updates and must + advance them using the original executor lifecycle before the next call. + """ + requests = batch.generation_requests + local_valid = ( + cell in self._cost_profile.cells + and not self._profile_call_active + and batch.num_context_requests == 0 + and batch.num_encoder_requests == 0 + and len(requests) == cell.graph_batch_size + and len({request.py_request_id for request in requests}) == len(requests) + and all( + not request.is_dummy + and request.py_beam_width == 1 + and not request.py_is_first_draft + and request.py_seq_slot is not None + for request in requests + ) + and self.runtime_draft_len == self._cost_profile.physical_k + and _valid_previous_state(new_tensors_device, requests, self._cost_profile.physical_k) + ) + row = [ + int(local_valid), + cell.graph_batch_size, + self._cost_profile.physical_k, + cell.verifier_budget, + len(self.cuda_graph_runner._profile_records), + ] + rows = self.dist.tp_allgather_int64(row).tolist() if self.mapping.tp_size > 1 else [row] + _require_shape_agreement(rows) + graph_runner = self.cuda_graph_runner + graph_runner.active_profile_cell = cell + graph_runner.active_request_ids = tuple(request.py_request_id for request in requests) + self._profile_call_active = True + try: + if cell.verifier_budget: + self._set_warmup_ragged_windows( + batch, cell.verifier_budget, self._cost_profile.physical_k + ) + else: + for request in requests: + request.py_verify_len = None + graph_runner.agreed_ragged_bucket = None + before = len(graph_runner._profile_records) + output = super().forward( + batch, + resource_manager=resource_manager, + new_tensors_device=new_tensors_device, + cache_indirection_buffer=cache_indirection_buffer, + ) + if len(graph_runner._profile_records) != before + 1: + raise RuntimeError( + "Profile forward did not execute exactly one original graph replay" + ) + return output + finally: + graph_runner.active_profile_cell = None + graph_runner.active_request_ids = () + self._profile_call_active = False + + def profile_observations(self) -> dict[str, Any]: + """Return raw observations, never an admitted table or enabled-performance result.""" + return { + "schema": "dspark-cost-profile-observation-v1", + "profile_only": True, + "physical_k": self._cost_profile.physical_k, + "record_confidence_head": self._cost_profile.record_confidence_head, + "timing_scope": "original-graph-replay-entry-including-static-copies", + "samples": self.cuda_graph_runner.completed_profile_samples(), + "full_policy_h": None, + "production_sps_table_eligible": False, + "production_admission_qualified": False, + } diff --git a/tensorrt_llm/_torch/pyexecutor/dspark_profile_graph_storage.py b/tensorrt_llm/_torch/pyexecutor/dspark_profile_graph_storage.py new file mode 100644 index 000000000000..607196ff9800 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/dspark_profile_graph_storage.py @@ -0,0 +1,188 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Offline-profiler metadata storage checks; no model or graph operations.""" + +from __future__ import annotations + +import json +from typing import Any + +import torch + +_V4_METADATA = ( + "tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.metadata", + "DeepseekV4TrtllmAttentionMetadata", +) +_TRT_METADATA_TYPES = { + _V4_METADATA, + ("tensorrt_llm._torch.attention.backends.sparse.dsa.metadata", "DSAtrtllmAttentionMetadata"), + ("tensorrt_llm._torch.attention.backends.trtllm", "TrtllmAttentionMetadata"), +} +_V4_BUFFER_MAPS = ( + "new_comp_kv_lens_cuda", + "cu_new_comp_kv_cuda", + "compressed_kv_lens_cuda", + "past_kv_lens_cuda", + "compressed_position_ids_cuda", + "compressed_mask_cuda", + "compress_block_tables", + "sparse_mla_topk_lens", +) + + +def _metadata_type(obj: Any) -> tuple[str, str]: + return type(obj).__module__, type(obj).__name__ + + +def _signature(value: torch.Tensor) -> dict[str, Any]: + return { + "address": value.data_ptr(), + "storage_offset": value.storage_offset(), + "shape": tuple(value.shape), + "stride": tuple(value.stride()), + "dtype": str(value.dtype), + "device": str(value.device), + } + + +def _role(obj: Any, name: str) -> str: + metadata_type = _metadata_type(obj) + if ( + name == "gen_indexer_kv_lens_cuda_runtime" + and metadata_type == _V4_METADATA + and obj.is_cuda_graph + and obj.num_contexts == 0 + and obj.num_generations > 0 + and obj._indexer_compress_ratio == 4 + ): + return "graph_produced_compressed_lengths" + if name == "workspace" and metadata_type in _TRT_METADATA_TYPES and obj.is_cuda_graph: + return "eager_only_workspace" + return "persistent" + + +def _refuse( + obj: Any, name: str, reason: str, before: Any, after: Any, context: dict[str, Any] | None = None +) -> None: + context = {} if context is None else context + details = { + "reason": reason, + "field": name, + "metadata_type": _metadata_type(obj), + "expected": before, + "actual": after, + "rank": context.get("rank"), + "graph_key": context.get("graph_key"), + } + raise RuntimeError( + "Profile graph storage contract failed: " + json.dumps(details, sort_keys=True) + ) + + +def _cuda_tensors(obj: Any, *, context: dict[str, Any] | None = None) -> dict[str, torch.Tensor]: + tensors = { + name: value + for name, value in vars(obj).items() + if isinstance(value, torch.Tensor) and value.is_cuda + } + if _metadata_type(obj) == _V4_METADATA: + for name in _V4_BUFFER_MAPS: + values = getattr(obj, name, None) + if not isinstance(values, dict) or any(type(key) is not int for key in values): + _refuse( + obj, + name, + "Invalid or missing required profile storage map", + "dict with integer keys", + type(values).__name__, + context, + ) + for key, value in values.items(): + if not isinstance(value, torch.Tensor) or not value.is_cuda: + _refuse( + obj, + f"{name}[{key}]", + "Invalid profile CUDA storage", + "CUDA tensor", + type(value).__name__, + context, + ) + tensors[f"{name}[{key}]"] = value + return tensors + + +def _validate_intermediate(obj: Any, value: torch.Tensor) -> None: + if ( + tuple(value.shape) != (obj.num_generations,) + or not value.is_contiguous() + or value.storage_offset() != 0 + or value.dtype != obj.kv_lens_cuda.dtype + or value.device != obj.kv_lens_cuda.device + ): + raise RuntimeError("Invalid profile graph-produced compressed-length tensor") + + +def snapshot_metadata_storage(obj: Any) -> dict[str, Any]: + """Record top-level CUDA tensors and audited V4 persistent buffer maps. + + This is a finite profiler contract, not an exhaustive traversal of arbitrary + runtime objects. Unknown tensors are persistent; no user-provided exemption + list or baseline reset is accepted. + """ + records = {} + for name, value in _cuda_tensors(obj).items(): + role = _role(obj, name) + if role == "graph_produced_compressed_lengths": + _validate_intermediate(obj, value) + records[name] = {"role": role, "signature": _signature(value)} + return {"metadata_type": _metadata_type(obj), "fields": records} + + +def validate_metadata_storage( + obj: Any, expected: dict[str, Any], *, context: dict[str, Any] | None = None +) -> None: + """Refuse genuine storage/layout drift with finite, field-specific details.""" + + def refuse(name: str, reason: str, before: Any, after: Any) -> None: + _refuse(obj, name, reason, before, after, context) + + if _metadata_type(obj) != tuple(expected["metadata_type"]): + refuse( + "", + "metadata type changed", + expected["metadata_type"], + _metadata_type(obj), + ) + actual = _cuda_tensors(obj, context=context) + before_fields = expected["fields"] + if set(actual) != set(before_fields): + changed = sorted(set(actual) ^ set(before_fields))[0] + refuse( + changed, + "CUDA field missing or introduced", + before_fields.get(changed), + _signature(actual[changed]) if changed in actual else None, + ) + for name, value in actual.items(): + before = before_fields[name] + role = _role(obj, name) + after = _signature(value) + if role != before["role"]: + refuse(name, "field role changed", before, {"role": role, "signature": after}) + if role == "graph_produced_compressed_lengths": + try: + _validate_intermediate(obj, value) + except RuntimeError as error: + refuse(name, str(error), before, after) + comparable_before = { + key: item for key, item in before["signature"].items() if key != "address" + } + comparable_after = {key: item for key, item in after.items() if key != "address"} + elif role == "eager_only_workspace": + # The graph uses cuda_graph_workspace, which remains persistent. + comparable_before = {key: before["signature"][key] for key in ("dtype", "device")} + comparable_after = {key: after[key] for key in ("dtype", "device")} + else: + comparable_before, comparable_after = before["signature"], after + if comparable_after != comparable_before: + refuse(name, "storage or layout changed", before, {"role": role, "signature": after}) diff --git a/tensorrt_llm/_torch/pyexecutor/dspark_request_profile.py b/tensorrt_llm/_torch/pyexecutor/dspark_request_profile.py new file mode 100644 index 000000000000..0784e84c8574 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/dspark_request_profile.py @@ -0,0 +1,535 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Explicit offline request ownership over the unchanged executor lifecycle.""" + +from __future__ import annotations + +import hashlib +import json +import math +from contextlib import contextmanager +from dataclasses import dataclass +from functools import partial +from typing import TYPE_CHECKING, Any, Optional + +from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator +from tensorrt_llm._torch.pyexecutor.dspark_profile_engine import ( + DSparkCostProfileEngine, + _valid_previous_state, +) +from tensorrt_llm._torch.pyexecutor.llm_request import ExecutorRequest, LlmRequestState +from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm._torch.pyexecutor.py_executor_creator import ( + _RuntimeConstructors, + create_py_executor, +) + +if TYPE_CHECKING: + from ray.actor import ActorHandle + +__all__ = ["OwnedDSparkProfileExecutor", "create_profile_executor"] + + +def _token_snapshot(request: Any) -> dict[str, Any]: + tokens = list(request.get_tokens(0)) + prompt = tokens[: request.py_orig_prompt_len] + return { + "request_id": request.py_request_id, + "process_local_request_object": id(request), + "seq_slot": request.py_seq_slot, + "batch_idx": request.py_batch_idx, + "decoding_iter": request.py_decoding_iter, + "tokens": tokens, + "original_input_sha256": hashlib.sha256( + json.dumps(prompt, separators=(",", ":")).encode() + ).hexdigest(), + "total_drafted": request.py_total_draft_tokens, + "total_accepted": request.py_total_accepted_draft_tokens, + } + + +@dataclass +class _ForwardJoin: + generation_id: int + executor_iteration: int + route: str + requests: tuple[Any, ...] + before: list[dict[str, Any]] + sample_index: int | None = None + update_before: list[dict[str, Any]] | None = None + sampler_update_applied: list[bool] | None = None + after: list[dict[str, Any]] | None = None + host_step_time_ms: float | None = None + + +class OwnedDSparkProfileEngine(DSparkCostProfileEngine): + """Only a source-owned executor boundary may enter a real request forward.""" + + def forward( + self, batch, resource_manager, new_tensors_device=None, cache_indirection_buffer=None + ): + if self.is_warmup: + return PyTorchModelEngine.forward( + self, + batch, + resource_manager, + new_tensors_device, + cache_indirection_buffer=cache_indirection_buffer, + ) + boundary = getattr(self, "_owned_forward_boundary", None) + if boundary is None or boundary[2] is not batch or boundary[3] is not new_tensors_device: + raise RuntimeError("No original-request owner boundary; ordinary serving denied") + route, cell, _, _ = boundary + if route == "cell": + return self.forward_profile_cell( + cell, + batch, + resource_manager, + new_tensors_device=new_tensors_device, + cache_indirection_buffer=cache_indirection_buffer, + ) + # Genuine prefill/first-draft/short tail/peer-not-ready transitions use + # the unchanged native path and produce no profile event record. + for request in batch.generation_requests: + request.py_verify_len = None + self.cuda_graph_runner.agreed_ragged_bucket = None + return PyTorchModelEngine.forward( + self, + batch, + resource_manager, + new_tensors_device, + cache_indirection_buffer=cache_indirection_buffer, + ) + + +class OwnedDSparkProfileExecutor(PyExecutor): + """Original executor loops plus fixed-cell provenance; not a serving factory.""" + + def __init__(self, *args, **kwargs): + self._owned_generation = None + self._owned_generation_ids = set() + self._owned_joins = [] + self._owned_pending = {} + self._owned_last_join = None + self._owned_request_objects = {} + self._owned_final_requests = {} + self._owned_failures = [] + self._owned_kv_capacity_enqueue = None + super().__init__(*args, **kwargs) + + @contextmanager + def _profiler(self): + """Join the original complete-loop timer, without new events or waits. + + The parent step reads the just-finished loop at the next loop boundary. + An overlap sampler may consume another batch there; its forward-event + duration and iterLatencyMS are not substitutes for this loop's cost. + """ + with super()._profiler() as step: + if self.is_warmup or self.model_engine._cost_profile.measurement_scope != "whole-step": + yield step + return + if not self.enable_iter_perf_stats: + raise ValueError("Whole-step costs require existing enable_iter_perf_stats") + previous_iteration = None + + def advance(): + nonlocal previous_iteration + step() + join = self._owned_joins[-1] if self._owned_joins else None + if ( + join is not None + and join.executor_iteration == previous_iteration + and join.route == "cell" + ): + duration = self._latest_host_step_time_ms + if ( + isinstance(duration, bool) + or not isinstance(duration, (int, float)) + or not math.isfinite(duration) + or duration <= 0 + ): + raise ValueError("Original cell loop has no finite positive hostStepTimeMS") + if join.host_step_time_ms is not None: + raise RuntimeError("Original cell loop timing was joined twice") + join.host_step_time_ms = duration + previous_iteration = self.iter_counter + + yield advance + + def begin_profile_generation(self, generation_id: int, cell: Any) -> None: + if type(generation_id) is not int or generation_id < 0: + raise ValueError("Original protocol generation ID must be a nonnegative integer") + self.model_engine._cost_profile.require_cell(cell) + with self.control_action(drain=True): + if self._owned_pending or generation_id in self._owned_generation_ids: + raise RuntimeError("Undrained or repeated original generation") + if self._owned_generation is not None: + raise RuntimeError("End the previous original generation before another") + self._owned_generation_ids.add(generation_id) + self._owned_generation = (generation_id, cell) + + def enqueue_requests(self, requests, result_wait_queue=None): + if self._owned_generation is None and not self.is_warmup: + authorization = self._owned_kv_capacity_enqueue + if ( + authorization is None + or requests is not authorization[0] + or authorization[1] + or result_wait_queue is not None + ): + raise RuntimeError("Ordinary requests require an explicit offline generation") + # The original KV estimator enqueues this exact internal list before + # setting warmup. Consume before delegation, even if enqueue fails. + self._owned_kv_capacity_enqueue = (requests, True) + return super().enqueue_requests(requests, result_wait_queue) + + def enqueue_request( + self, request: ExecutorRequest, result_wait_queue: "Optional[ActorHandle]" = None + ) -> int: + if self._owned_generation is None and not self.is_warmup: + raise RuntimeError("Ordinary requests require an explicit offline generation") + return super().enqueue_request(request, result_wait_queue) + + def _handle_errors( + self, + error_msg=None, + *, + requests=None, + charge_budget=True, + fatal_is_collective_aligned=False, + ): + if self._owned_generation is not None: + self._owned_failures.append(str(error_msg or "Original executor error")) + return super()._handle_errors( + error_msg, + requests=requests, + charge_budget=charge_budget, + fatal_is_collective_aligned=fatal_is_collective_aligned, + ) + + def _forward_step(self, batch, new_tensors_device=None): + if self.is_warmup: + return super()._forward_step(batch, new_tensors_device) + generation = self._owned_generation + engine = self.model_engine + requests = tuple(request for request in batch.all_requests() if not request.is_dummy) + ready = generation is not None + generation_id, cell = generation if ready else (-1, None) + eligible = ( + ready + and batch.num_context_requests == 0 + and batch.num_encoder_requests == 0 + and len(batch.generation_requests) == cell.graph_batch_size + and all( + not r.is_dummy + and not r.py_is_first_draft + and r.py_beam_width == 1 + and r.py_seq_slot is not None + for r in batch.generation_requests + ) + and len({r.py_request_id for r in batch.generation_requests}) + == len(batch.generation_requests) + and _valid_previous_state( + new_tensors_device, batch.generation_requests, engine._cost_profile.physical_k + ) + ) + row = [ + int(ready), + generation_id, + self.iter_counter, + cell.graph_batch_size if ready else -1, + cell.verifier_budget if ready else -1, + int(eligible), + ] + rows = self.dist.tp_allgather_int64(row).tolist() if self.dist.tp_size > 1 else [row] + if ( + len(rows) != self.dist.tp_size + or any(len(r) != 6 or r[0] != 1 for r in rows) + or any(r[:5] != rows[0][:5] for r in rows) + ): + raise RuntimeError("All ranks must agree on the original generation/iteration/cell") + route = "cell" if all(r[5] == 1 for r in rows) else "native-transition" + for request in requests: + key = (generation_id, request.py_request_id) + previous = self._owned_request_objects.setdefault(key, request) + if previous is not request: + raise RuntimeError("Request ID reused with a different original incarnation") + join = _ForwardJoin( + generation_id, + self.iter_counter, + route, + requests, + [_token_snapshot(request) for request in requests], + ) + before_records = len(engine.cuda_graph_runner._profile_records) + if getattr(engine, "_owned_forward_boundary", None) is not None: + raise RuntimeError("Overlapping owner forward boundaries") + engine._owned_forward_boundary = (route, cell, batch, new_tensors_device) + try: + output = super()._forward_step(batch, new_tensors_device) + if output is None: + raise RuntimeError("Original forward failed; no observation promotion") + count = len(engine.cuda_graph_runner._profile_records) + if count != before_records + int(route == "cell"): + raise RuntimeError("Native transition or fixed-cell replay count differs") + if route == "cell": + record = engine.cuda_graph_runner._profile_records[-1] + if tuple(record["request_ids"]) != tuple( + r.py_request_id for r in batch.generation_requests + ): + raise RuntimeError("Captured replay and original request identities differ") + join.sample_index = record["sample_index"] + self._owned_joins.append(join) + self._owned_last_join = join + return output + except BaseException as error: + self._owned_failures.append(type(error).__name__ + ": " + str(error)) + raise + finally: + engine._owned_forward_boundary = None + + def _sample_async(self, batch, output): + state = super()._sample_async(batch, output) + if self.is_warmup: + return state + join = self._owned_last_join + if state is None or join is None or id(state) in self._owned_pending: + raise RuntimeError("Missing or reused original sampler state") + if any( + not any(request is original for original in join.requests) + for request in state.requests + if not request.is_dummy + ): + raise RuntimeError("Sampler state belongs to another original forward") + if join.sample_index is not None and tuple(state.requests) != join.requests: + raise RuntimeError("Measured sampler requests differ from original decode forward") + self._owned_pending[id(state)] = (state, join) + self._owned_last_join = None + return state + + def _update_requests(self, state, resource_manager=None): + if self.is_warmup: + return super()._update_requests(state, resource_manager) + pending = self._owned_pending.get(id(state)) + if pending is None or pending[0] is not state: + raise RuntimeError("Unknown or already consumed original sampler state") + join = pending[1] + join.update_before = [_token_snapshot(request) for request in join.requests] + completed_before_update = [ + request.state == LlmRequestState.GENERATION_COMPLETE for request in join.requests + ] + super()._update_requests(state, resource_manager) + if join.sample_index is not None: + self.model_engine.cuda_graph_runner.complete_sample(join.sample_index) + # Read only after the original sampler sync/token/KV/stop update and + # original accumulation. Under overlap this is the previous iteration. + join.after = [_token_snapshot(request) for request in join.requests] + join.sampler_update_applied = [ + after["decoding_iter"] > before["decoding_iter"] + for before, after in zip(join.update_before, join.after) + ] + if join.sample_index is not None and any( + not applied and not completed + for applied, completed in zip(join.sampler_update_applied, completed_before_update) + ): + self._owned_failures.append("Original measured request update did not advance") + raise RuntimeError("Original measured request update did not advance") + del self._owned_pending[id(state)] + + def _free_request_resources(self, request): + generation = self._owned_generation + final = None + if generation is not None and not request.is_dummy: + key = (generation[0], request.py_request_id) + if ( + self._owned_request_objects.get(key) is request + and request.state == LlmRequestState.GENERATION_COMPLETE + ): + final = (key, _token_snapshot(request)) + result = super()._free_request_resources(request) + if final is not None: + self._owned_final_requests[final[0]] = final[1] + return result + + def _profile_drain_predicates(self): + previous = self.previous_batch + state = previous.sample_state if previous is not None else None + pending = self._owned_pending.get(id(state)) if state is not None else None + join = pending[1] if pending is not None else None + return { + "generation_id": self._owned_generation[0] if self._owned_generation else None, + "failures": len(self._owned_failures), + "active_requests": len(self.active_requests), + "waiting_requests": len(self.waiting_queue), + "pending_samples": len(self._owned_pending), + "last_join_present": self._owned_last_join is not None, + "forward_boundary_present": ( + getattr(self.model_engine, "_owned_forward_boundary", None) is not None + ), + "previous_batch_present": previous is not None, + "pending_matches_previous": ( + state is not None + and len(self._owned_pending) == 1 + and pending is not None + and pending[0] is state + and any(join is original for original in self._owned_joins) + and self._owned_generation is not None + and join.generation_id == self._owned_generation[0] + and join.update_before is None + and join.sampler_update_applied is None + and join.after is None + ), + } + + def end_profile_generation(self, *, control=None) -> dict[str, Any]: + if self._owned_generation is None: + raise RuntimeError("No original generation to drain") + with self.control_action(drain=True): + before = self._profile_drain_predicates() + if control is None: + raise RuntimeError( + "Offline profile drain needs the original gang-control communicator" + ) + peers = control.allgather(before) + quiescent = all( + row["failures"] == 0 + and row["active_requests"] == 0 + and row["waiting_requests"] == 0 + and not row["last_join_present"] + and not row["forward_boundary_present"] + and row["generation_id"] == before["generation_id"] + for row in peers + ) + no_tail = all(row["pending_samples"] == 0 for row in peers) + exact_tail = all( + row["pending_samples"] == 1 and row["pending_matches_previous"] for row in peers + ) + if not quiescent or not (no_tail or exact_tail): + raise RuntimeError( + f"Original generation has not cleanly drained: gang_predicates={peers!r}" + ) + consumed = False + if exact_tail: + # All ranks are paused in the original exclusive control action. + # Delegate every sampler/update/resource/timing operation to the + # existing executor helper; do not clear pending joins directly. + self._consume_previous_batch_for_rebalance() + consumed = True + after = self._profile_drain_predicates() + if ( + self._owned_pending + or self._owned_last_join is not None + or self._owned_failures + or self.active_requests + or len(self.waiting_queue) + ): + raise RuntimeError( + "Original generation has not cleanly drained: " + f"before={before!r}, after={after!r}" + ) + generation_id = self._owned_generation[0] + keys = [key for key in self._owned_request_objects if key[0] == generation_id] + if not keys or any(key not in self._owned_final_requests for key in keys): + raise RuntimeError("Original request final resource-release joins are incomplete") + observation = self.model_engine.profile_observations() + by_sample = { + join.sample_index: join + for join in self._owned_joins + if join.sample_index is not None + } + if set(by_sample) != {row["sample_index"] for row in observation["samples"]}: + raise RuntimeError("Completed replay events and lifecycle joins differ") + joins = [ + { + "generation_id": join.generation_id, + "executor_iteration": join.executor_iteration, + "route": join.route, + "sample_index": join.sample_index, + "before": join.before, + "update_before": join.update_before, + "sampler_update_applied": join.sampler_update_applied, + "after": join.after, + "host_step_time_ms": join.host_step_time_ms, + } + for join in self._owned_joins + if join.generation_id == generation_id + ] + if any(join["after"] is None for join in joins): + raise RuntimeError("Original sampling/update join missing") + if self.model_engine._cost_profile.measurement_scope == "whole-step" and any( + join["sample_index"] is not None and join["host_step_time_ms"] is None + for join in joins + ): + raise RuntimeError("Measured cell lacks its original complete-loop timing") + generation_samples = { + join["sample_index"] for join in joins if join["sample_index"] is not None + } + if not generation_samples: + raise RuntimeError("Prescribed generation did not execute a measured graph cell") + observation = { + **observation, + "samples": [ + row + for row in observation["samples"] + if row["sample_index"] in generation_samples + ], + } + self._owned_generation = None + return { + "schema": "dspark-original-request-profile-generation-v1", + "generation_id": generation_id, + "drain": { + "before": before, + "after": after, + "original_previous_batch_consumed": consumed, + }, + "observation": observation, + "lifecycle_joins": joins, + "original_final_requests": [self._owned_final_requests[key] for key in keys], + "original_shutdown_still_required": True, + "task_accuracy_qualified": False, + "native_or_gpu_qualified": False, + "full_policy_h": None, + "production_admission_qualified": False, + } + + +class _ProfileKvCacheCreator(KvCacheCreator): + def configure_kv_cache_capacity(self, py_executor=None) -> None: + if py_executor is None or self._skip_est: + return super().configure_kv_cache_capacity(py_executor) + if ( + not isinstance(py_executor, OwnedDSparkProfileExecutor) + or py_executor._owned_generation is not None + or py_executor._owned_kv_capacity_enqueue is not None + ): + raise RuntimeError("Invalid or nested offline KV startup phase") + py_executor._owned_kv_capacity_enqueue = (self._dummy_reqs, False) + try: + return super().configure_kv_cache_capacity(py_executor) + finally: + py_executor._owned_kv_capacity_enqueue = None + + +def create_profile_executor(*, llm_args, profile, checkpoint_dir=None, tokenizer=None): + """Construct only the prescribed offline owner; never a serving executor. + + No source pins, private manifests, historical receipts or global cloning are + required. Runtime and loaded-byte identity is measured by the public driver. + Normal factory construction, resource/KV estimation and loops are retained. + """ + if llm_args.enable_resource_governor: + raise ValueError("The offline profile owner does not support resource-governor control") + constructors = _RuntimeConstructors( + model_engine=partial(OwnedDSparkProfileEngine, profile=profile), + executor=OwnedDSparkProfileExecutor, + kv_cache_creator=_ProfileKvCacheCreator, + ) + return create_py_executor( + llm_args=llm_args, + checkpoint_dir=checkpoint_dir, + tokenizer=tokenizer, + runtime_constructors=constructors, + ) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/metadata.py b/tensorrt_llm/_torch/pyexecutor/engine/metadata.py index c9ed43c82ef9..9cd712e5ca85 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/metadata.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/metadata.py @@ -53,6 +53,7 @@ def build_attention_metadata( draft_kv_cache_manager: KVCacheManager | KVCacheManagerV2 | None = None, enable_context_mla_with_cached_kv: bool | None = None, num_heads_per_kv: int | None = None, + enable_ragged_verification: bool = False, ) -> AttentionMetadata: """Construct attention metadata and resolve model-derived inputs.""" pretrained_config = model_config.pretrained_config @@ -69,6 +70,15 @@ def build_attention_metadata( if sparse_attention_config is not None else None ) + metadata_kwargs = {} + if enable_ragged_verification: + from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import ( + DSAtrtllmAttentionMetadata, + ) + + if not issubclass(attention_backend.Metadata, DSAtrtllmAttentionMetadata): + raise ValueError("Ragged verification requires the DSA target metadata backend") + metadata_kwargs["enable_ragged_verification"] = True return attention_backend.Metadata( max_num_requests=max_batch_size, max_num_tokens=max_num_tokens, @@ -82,4 +92,5 @@ def build_attention_metadata( cache_indirection=cache_indirection, num_heads_per_kv=num_heads_per_kv, sparse_metadata_params=sparse_metadata_params, + **metadata_kwargs, ) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index c74a0a8d9e7e..ded8e8c51cd6 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -11,7 +11,7 @@ import os from collections.abc import Callable, Iterator, Sequence from contextlib import contextmanager -from typing import Any +from typing import Any, List, Optional, Tuple import torch @@ -38,6 +38,7 @@ from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( CUDAGraphRunner, CUDAGraphRunnerConfig, + DSparkCUDAGraphRunner, KeyType, get_mrope_dummy_seq_slot, ) @@ -47,7 +48,11 @@ BaseMambaCacheManager, MambaHybridCacheManager, ) -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, get_draft_token_length +from tensorrt_llm._torch.pyexecutor.llm_request import ( + LlmRequest, + get_draft_token_length, + get_request_tokens_per_gen_step, +) from tensorrt_llm._torch.pyexecutor.resource_manager import ( BaseResourceManager, KVCacheManager, @@ -71,6 +76,7 @@ ) from tensorrt_llm._torch.speculative.dspark import DSparkSpecMetadata from tensorrt_llm._torch.speculative.interface import INVALID_PROMPT_LOOKAHEAD_TOKEN +from tensorrt_llm._torch.speculative.ragged_helpers import ragged_gather_index_lists from tensorrt_llm._torch.speculative.spec_sampler_base import SampleStateTensorsSpec from tensorrt_llm._torch.speculative.utils import get_static_draft_len, resolve_draft_len from tensorrt_llm._torch.utils import ( @@ -485,6 +491,14 @@ def _allocate_decoder_buffers(self) -> None: self.previous_batch_indices_cuda = torch.empty( (self._config.max_num_tokens,), dtype=torch.int, device="cuda" ) + self.ragged_verify_lens_cuda = torch.empty( + (self._config.max_batch_size + 1,), dtype=torch.int, device="cuda" + ) + self.ragged_qo_indptr_cuda = torch.empty( + (self._config.max_batch_size + 2,), dtype=torch.int, device="cuda" + ) + self._pinned_host_cache = {} + self._pinned_host_active = {} self.input_ids_cuda = torch.empty( (self._config.max_num_tokens,), dtype=torch.int, device="cuda" ) @@ -512,6 +526,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_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 # Token shape of the last autotuner warmup that ran a forward pass. self._completed_autotuner_warmup_num_tokens: int | None = None @@ -542,12 +629,17 @@ def _init_decoder_state(self) -> None: self._dynamic_draft_len_mapping = self._compute_dynamic_draft_len_mapping() def _initialize_cuda_graph_runner(self) -> CUDAGraphRunner | None: - cuda_graph_runner = CUDAGraphRunner(self._cuda_graph_runner_config()) + cuda_graph_runner = self._create_cuda_graph_runner(self._cuda_graph_runner_config()) register_hooks = getattr(self.model, "register_cuda_graph_pre_replay_hooks", None) if register_hooks is not None: register_hooks(cuda_graph_runner) return cuda_graph_runner + def _create_cuda_graph_runner(self, config: CUDAGraphRunnerConfig) -> CUDAGraphRunner: + """Construction seam for explicit offline tooling, with ordinary defaults retained.""" + runner_type = DSparkCUDAGraphRunner if self._dspark_confidence_enabled else CUDAGraphRunner + return runner_type(config) + def _cuda_graph_runner_config(self) -> CUDAGraphRunnerConfig: return CUDAGraphRunnerConfig( use_cuda_graph=( @@ -773,7 +865,7 @@ def _warmup_scheduled(self, resource_manager: ResourceManager, kv_cache_manager) # The lifetime of model engine and kv cache manager can be different. # Reset the global cuda graph dummy requests in warmup. - self.cuda_graph_runner.padding_dummy_requests = {} + self.cuda_graph_runner.reset_padding_dummies() if self.mapping.cp_size > 1: cp_type = self.mapping.cp_config.get("cp_type", None) @@ -881,6 +973,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") @@ -1956,7 +2050,9 @@ def _compute_dynamic_draft_len_mapping(self) -> dict[int, int] | None: mapping[graph_bs] = draft_len return mapping - def _get_graphs_to_capture(self, cuda_graph_batch_sizes: list[int]) -> list[tuple[int, int]]: + def _get_graphs_to_capture( + self, cuda_graph_batch_sizes: list[int] + ) -> list[tuple[int, int] | tuple[int, int, int]]: """Determine which (batch_size, draft_len) graphs to capture. Returns: @@ -1996,6 +2092,33 @@ def _get_graphs_to_capture(self, cuda_graph_batch_sizes: list[int]) -> list[tupl ) return graphs + if self._dspark_confidence_enabled: + max_draft_len = int(self._config.spec_config.max_draft_len) + native_graphs = [(bs, max_draft_len) for bs in cuda_graph_batch_sizes] + if not self._dspark_trims_submitted_tokens: + return native_graphs + from tensorrt_llm._torch.speculative.dspark_planner import ExactSpsCostTable + + exact_table = self._dspark_sps_cost_table + if not isinstance(exact_table, ExactSpsCostTable): + raise RuntimeError( + "DSpark confidence graph capture requires an authenticated exact SPS table" + ) + ragged_graphs = [ + (bs, max_draft_len, verifier_budget) + for bs in cuda_graph_batch_sizes + for verifier_budget in exact_table.production_candidate_budgets(bs) + ] + # Full K keeps the ordinary native graph. Compact layouts have + # distinct keys; they cannot alias its uniform attention layout. + graphs = native_graphs + ragged_graphs + logger.info( + f"DSpark ragged verification: capturing {len(graphs)} graphs " + f"({len(native_graphs)} native static-K{max_draft_len} graphs + " + f"{len(ragged_graphs)} exact measured (G,V) cells)." + ) + return graphs + # Case 2: Static draft length # Match the runtime_draft_len semantics enforced in _prepare_tp_inputs: # logical K for linear-tree modes, total tree tokens for tree decoding. @@ -2025,7 +2148,7 @@ def _warmup_and_capture_cuda_graphs(self, resource_manager: ResourceManager) -> self._run_cuda_graph_warmup(resource_manager) finally: self.cuda_graph_runner.is_warmup_only = False - self.cuda_graph_runner.padding_dummy_requests = {} + self.cuda_graph_runner.reset_padding_dummies() # Captured graphs keep the MoE communication timeouts that were # current at capture, and serving replays them, so capture with the # serving timeouts. The warmup-only pass above keeps the warmup ones. @@ -2162,7 +2285,9 @@ def _run_capture_pass( pinned_tier = sample_type or SampleType.FULL self.cuda_graph_runner.set_capture_sample_type(pinned_tier) try: - for bs, draft_len in graphs_to_capture: + for entry in graphs_to_capture: + bs, draft_len = entry[:2] + verify_bucket = entry[2] if len(entry) == 3 else None if bs > self._config.max_batch_size: continue @@ -2175,6 +2300,8 @@ def _run_capture_pass( force_non_greedy=force_non_greedy, ) with self._release_batch_context(warmup_request, resource_manager) as batch: + if batch is not None and verify_bucket is not None: + self._set_warmup_ragged_windows(batch, verify_bucket, draft_len) if batch is None: # No KV cache space for this batch size. During KV # cache estimation this makes the profiling peak @@ -2211,6 +2338,8 @@ def _run_capture_pass( finally: self._force_lora_graph_for_capture = None self.cuda_graph_runner.set_capture_sample_type(None) + if self._dspark_confidence_enabled: + self.cuda_graph_runner.agreed_ragged_bucket = None if self._lora.cuda_graph_manager is None: lora_graph_cases = [False] @@ -2261,7 +2390,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 @@ -2790,6 +2920,11 @@ def _add_dummy_request_resources( """Allocate resources dummy requests need beyond the KV caches.""" return True + @property + def _ragged_verification_enabled(self) -> bool: + """Select persistent ragged storage before attention metadata construction.""" + return self._dspark_confidence_enabled + def _set_up_attn_metadata( self, kv_cache_manager: KVCacheManager | KVCacheManagerV2, @@ -2815,6 +2950,7 @@ def _set_up_attn_metadata( else None, kv_cache_manager=kv_cache_manager, draft_kv_cache_manager=draft_kv_cache_manager, + enable_ragged_verification=self._ragged_verification_enabled, ) if isinstance(kv_cache_manager, BaseMambaCacheManager): self.attn_metadata.mamba_chunk_size = getattr( @@ -3048,7 +3184,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 @@ -3057,19 +3195,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. @@ -3130,6 +3289,828 @@ 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", + "batch_size_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 = [] + original_batch_size = metadata.batch_size_cuda.clone() + 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.batch_size_cuda.fill_(num_generations) + metadata._compute_compressed_mask( + metadata.new_comp_kv_lens_cuda, + metadata.cu_new_comp_kv_cuda, + metadata.compressed_mask_cuda, + metadata.batch_size_cuda, + compress_ratios, + ) + torch.cuda.synchronize() + finally: + metadata.batch_size_cuda.copy_(original_batch_size) + for ratio in prepared_ratios: + total_tokens = total_compressed_tokens[ratio] + metadata.past_kv_lens_cuda[ratio][:num_generations].zero_() + 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_schedule import DSparkFusedScheduleError + from tensorrt_llm._torch.speculative.dspark_select import ( + gather_packed_draft_tokens, + materialize_compact_layout, + select_windows_device, + ) + + 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() + if worker is None: + return False + planner = worker.verify_planner + if planner is None or worker.staged_confidence_buffer() is None: + return False + if worker.batch_slot_view(1) is None: + return False + spec_metadata = inputs.get("spec_metadata") + attn_metadata = inputs.get("attn_metadata") + if spec_metadata is None or attn_metadata is None: + return False + if self._config.use_mrope: + # Host-selected ragged windows build the complete three-axis + # MRoPE positions correctly. The device prologue currently owns + # only the scalar position buffer, so decline fresh reranking + # rather than leave the 3-D positions on the stale shape split. + return False + apply_device_layout = getattr(attn_metadata, "apply_device_ragged_layout", None) + if not callable(apply_device_layout): + # Device-selected row ownership is an attention-backend + # capability. Decline before mutating any shared buffers when a + # backend only supports the ordinary host-selected ragged layout. + return False + if int(attn_metadata.num_contexts) != 0: + return False + + bucket = int(runner.agreed_ragged_bucket) + padded_bs = int(self._dspark_last_padded_bs) + n_real = int(self._dspark_last_num_real) + pad_len_tok = int(runner.ragged_pad_verify_len) + 1 + real_tokens = bucket - (padded_bs - n_real) * pad_len_tok + cfg = planner.cfg + # The scheduler cannot grant more than the real rows can absorb under + # the published split; the fill tops up any shortfall. + budget = max(0, min(int(budget), real_tokens - n_real * (cfg.min_verify_len + 1))) + + 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, @@ -3673,6 +4654,9 @@ def _helix_pack_extend(request, group: int) -> int: # attention-DP dummies) whose input ids are not appended to the host # ``input_ids`` list; zero-filled after the input-id copies below. unwritten_input_id_rows = [] + previous_batch_token_start = None + previous_batch_tokens_per_request = [] if self._dspark_trims_submitted_tokens else None + extend_tokens_emitted = 0 runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step(runtime_draft_len) runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 for request in extend_requests: @@ -3692,6 +4676,11 @@ def _helix_pack_extend(request, group: int) -> int: or request.is_dummy or request.py_batch_idx is None ): + num_draft_tokens = get_draft_token_length(request) + if self._dspark_trims_submitted_tokens: + num_draft_tokens = ( + get_request_tokens_per_gen_step(request, 1 + num_draft_tokens) - 1 + ) # get token ids, including input token ids and draft token ids. For these dummy requests, # no need to copy the token ids. if not (request.is_attention_dp_dummy or request.is_cuda_graph_dummy): @@ -3699,10 +4688,12 @@ def _helix_pack_extend(request, group: int) -> int: input_ids.append(request.get_tokens(0)[request.context_current_position]) else: input_ids.append(request.get_last_tokens(0)) - input_ids.extend(request.py_draft_tokens) - draft_tokens.extend(request.py_draft_tokens) + submitted_draft_tokens = request.py_draft_tokens + if self._dspark_trims_submitted_tokens: + submitted_draft_tokens = submitted_draft_tokens[:num_draft_tokens] + input_ids.extend(submitted_draft_tokens) + draft_tokens.extend(submitted_draft_tokens) # get other ids and lengths - num_draft_tokens = get_draft_token_length(request) past_seen_token_num = ( request.context_current_position if is_promoted_context @@ -3737,6 +4728,9 @@ def _helix_pack_extend(request, group: int) -> int: num_cached_tokens_per_seq.append( past_seen_token_num - request.py_num_compressed_tokens ) + if self._dspark_trims_submitted_tokens: + request.cached_tokens = past_seen_token_num + extend_tokens_emitted += 1 + num_draft_tokens if _has_cp_helix: # Verify group [base, base+group) in GLOBAL positions. # On a helix gen worker the request's token list is the @@ -3760,30 +4754,36 @@ def _helix_pack_extend(request, group: int) -> int: previous_batch_idx = request.py_batch_idx request.py_batch_idx = request.py_seq_slot - sequence_lengths.append(runtime_tokens_per_gen_step) + request_tokens_per_step = runtime_tokens_per_gen_step + if self._dspark_trims_submitted_tokens: + request_tokens_per_step = get_request_tokens_per_gen_step( + request, runtime_tokens_per_gen_step + ) + sequence_lengths.append(request_tokens_per_step) num_accepted_draft_tokens.append(request.py_num_accepted_draft_tokens) past_seen_token_num = request.max_beam_num_tokens - 1 - draft_lens.append(runtime_draft_token_buffer_width) + draft_lens.append(request_tokens_per_step - 1) gather_ids.extend( - list(range(len(position_ids), len(position_ids) + runtime_tokens_per_gen_step)) + list(range(len(position_ids), len(position_ids) + request_tokens_per_step)) ) position_ids.extend( - list( - range( - past_seen_token_num, past_seen_token_num + runtime_tokens_per_gen_step - ) - ) + list(range(past_seen_token_num, past_seen_token_num + request_tokens_per_step)) ) # previous tensor previous_batch_indices.append(previous_batch_idx) - previous_pos_indices.extend([previous_batch_idx] * runtime_tokens_per_gen_step) + previous_pos_indices.extend([previous_batch_idx] * request_tokens_per_step) + if previous_batch_token_start is None: + previous_batch_token_start = extend_tokens_emitted + if previous_batch_tokens_per_request is not None: + previous_batch_tokens_per_request.append(request_tokens_per_step) + extend_tokens_emitted += request_tokens_per_step num_cached_tokens_per_seq.append( - past_seen_token_num - + runtime_tokens_per_gen_step - - request.py_num_compressed_tokens + past_seen_token_num + request_tokens_per_step - request.py_num_compressed_tokens ) + if self._dspark_trims_submitted_tokens: + request.cached_tokens = past_seen_token_num + request_tokens_per_step if _has_cp_helix: # In-flight predecessor: mirror the non-helix convention # above -- positions are packed from the stale base (the @@ -3792,7 +4792,7 @@ def _helix_pack_extend(request, group: int) -> int: # in recompute_helix_spec_buffers overrides them). The # base is reconstructed GLOBALLY (see the no-previous # branch: the token list is rank-local under helix). - group = runtime_tokens_per_gen_step + group = request_tokens_per_step base = _helix_pack_extend(request, group) local_full = _helix_local_len_host(base + group) helix_owned_new_tokens.append(0) @@ -3802,7 +4802,7 @@ def _helix_pack_extend(request, group: int) -> int: and spec_config.spec_dec_mode.extend_ctx(self._config.attention_backend) and spec_config.is_linear_tree ): - prompt_lengths.append(runtime_tokens_per_gen_step) + prompt_lengths.append(request_tokens_per_step) else: prompt_lengths.append(request.py_prompt_len) @@ -4056,6 +5056,9 @@ def _helix_pack_extend(request, group: int) -> int: mrope_delta_read_seq_slots.clear() previous_batch_len = len(previous_batch_indices) + self._dspark_prev_covers_batch = previous_batch_len > 0 and previous_batch_len == len( + extend_requests + ) - len(extend_dummy_requests) def previous_seq_slots_device(): previous_batch_indices_host = torch.tensor( @@ -4099,18 +5102,43 @@ def previous_seq_slots_device(): if previous_batch_len > 0: previous_slots = previous_seq_slots_device() + is_ragged_gen = previous_batch_tokens_per_request is not None and any( + tokens != runtime_tokens_per_gen_step + for tokens in previous_batch_tokens_per_request + ) # previous input ids - previous_batch_tokens = previous_batch_len * runtime_tokens_per_gen_step - new_tokens = new_tokens_device.transpose(0, 1)[ - previous_slots, :runtime_tokens_per_gen_step - ].flatten() + previous_batch_tokens = len(previous_pos_indices) + if is_ragged_gen: + token_rows, token_cols = self._ragged_gather_indices( + previous_batch_indices, previous_batch_tokens_per_request + ) + new_tokens = new_tokens_device.transpose(0, 1)[token_rows, token_cols].flatten() + else: + new_tokens = new_tokens_device.transpose(0, 1)[ + previous_slots, :runtime_tokens_per_gen_step + ].flatten() self.input_ids_cuda[num_tokens : num_tokens + previous_batch_tokens].copy_( new_tokens, non_blocking=True ) # previous draft tokens - previous_batch_draft_tokens = previous_batch_len * runtime_draft_token_buffer_width - if runtime_draft_token_buffer_width > 0: + if is_ragged_gen: + previous_batch_draft_tokens = previous_batch_tokens - previous_batch_len + if previous_batch_draft_tokens > 0: + draft_rows, draft_cols = self._ragged_gather_indices( + previous_batch_indices, + [tokens - 1 for tokens in previous_batch_tokens_per_request], + ) + self.draft_tokens_cuda[ + num_draft_tokens : num_draft_tokens + previous_batch_draft_tokens + ].copy_( + next_draft_tokens_device[draft_rows, draft_cols].flatten(), + non_blocking=True, + ) + elif runtime_draft_token_buffer_width > 0: + previous_batch_draft_tokens = ( + previous_batch_len * runtime_draft_token_buffer_width + ) self.draft_tokens_cuda[ num_draft_tokens : num_draft_tokens + previous_batch_draft_tokens ].copy_( @@ -4120,7 +5148,18 @@ def previous_seq_slots_device(): non_blocking=True, ) # prepare data for the preprocess inputs - kv_len_offsets_device = new_tokens_lens_device - runtime_tokens_per_gen_step + if is_ragged_gen: + tokens_per_previous_request = torch.tensor( + previous_batch_tokens_per_request, + dtype=torch.long, + pin_memory=prefer_pinned(), + ).to(new_tokens_lens_device.device, non_blocking=True) + previous_kv_len_offsets = ( + new_tokens_lens_device[previous_slots] - tokens_per_previous_request + ).clamp_(min=-(runtime_draft_len + 1), max=runtime_draft_len + 1) + else: + kv_len_offsets_device = new_tokens_lens_device - runtime_tokens_per_gen_step + previous_kv_len_offsets = kv_len_offsets_device[previous_slots] previous_pos_indices_host = torch.tensor( previous_pos_indices, dtype=torch.int, pin_memory=prefer_pinned() ) @@ -4149,10 +5188,14 @@ def previous_seq_slots_device(): # Already set to '0' during initialization. num_extend_reqeust_wo_dummy = len(extend_requests) - len(extend_dummy_requests) - self.previous_pos_id_offsets_cuda[ - (num_extend_reqeust_wo_dummy - previous_batch_len) - * runtime_tokens_per_gen_step : num_extend_reqeust_wo_dummy + pos_offsets_start = ( + previous_batch_token_start + if previous_batch_token_start is not None + else (num_extend_reqeust_wo_dummy - previous_batch_len) * runtime_tokens_per_gen_step + ) + self.previous_pos_id_offsets_cuda[ + pos_offsets_start : pos_offsets_start + previous_batch_tokens ].copy_( new_tokens_lens_device[self.previous_pos_indices_cuda[0:previous_batch_tokens]], non_blocking=True, @@ -4160,7 +5203,7 @@ def previous_seq_slots_device(): self.previous_kv_lens_offsets_cuda[ num_extend_reqeust_wo_dummy - previous_batch_len : num_extend_reqeust_wo_dummy - ].copy_(kv_len_offsets_device[previous_slots], non_blocking=True) + ].copy_(previous_kv_len_offsets, non_blocking=True) elif new_tokens_device is not None: seq_slots_device = previous_seq_slots_device() @@ -4254,7 +5297,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( @@ -4484,6 +5542,10 @@ def previous_seq_slots_device(): ] if context_prompt_lookahead is not None: spec_metadata.populate_context_prompt_lookahead(context_prompt_lookahead) + if self._dspark_trims_submitted_tokens: + self._attach_ragged_verify_layout( + spec_metadata, attn_metadata, scheduled_requests.generation_requests + ) # No-op for non 1-model spec_metadata.populate_sampling_params_for_one_model(scheduled_requests.all_requests()) spec_metadata.prepare() @@ -4528,6 +5590,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 @@ -4787,6 +5856,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: @@ -4890,6 +5960,13 @@ def _forward_decoder( self._prepare_inputs_event = torch.cuda.Event() self._prepare_inputs_event.record() + if ( + can_run_graph + and not is_dummy + and getattr(self, "_dspark_device_budget", None) is not None + ): + self._apply_device_window_prologue(inputs, new_tensors_device) + breakable_runner = self.breakable_cuda_graph_runner with with_shared_pool(self.cuda_graph_runner.get_graph_pool()): @@ -5145,6 +6222,16 @@ def _execute_logit_post_processors(self, scheduled_requests: ScheduledRequests, logits_processors = getattr(request, "py_logits_post_processors", None) if logits_processors: + if ( + not is_context_request + and getattr(request, "py_verify_len", None) is not None + ): + raise RuntimeError( + "per-request logits post-processors are not supported with DSpark " + "ragged verification: the processor interface assumes one uniform " + "row stride and would modify another request's packed logits after " + "a short verify window" + ) token_ids = ( [request.get_tokens(0)] if is_context_request diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 221f9145b17c..c14b982fefdd 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -804,6 +804,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 c51ec2bc75e7..033fcd972890 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -79,6 +79,7 @@ from .adp_iter_stats import ADPIterStatsBuffer from .connectors.kv_cache_connector import KvCacheConnectorManager from .connectors.kv_cache_layout import build_kv_cache_layout_v2 +from .cuda_graph_runner import ADPShapeAgreement from .disagg_adapter import PyExecutorEffects, PyExecutorRequestRegistry from .dwdp import DwdpManager from .error_classification import ErrorBudget @@ -133,6 +134,150 @@ _UNBOUNDED_STATS_MAX_LEN = -1 _V2_RETENTION_IGNORED_LOG_KEY = "kv_cache_manager_v2_retention_config_ignored" +_DSPARK_EXACT_WIRE_PREFIX_LEN = 16 +_DSPARK_ADP_WIRE_TRAILER_LEN = 3 +_DSPARK_EXACT_NATIVE_YIELD_INDEX = _DSPARK_EXACT_WIRE_PREFIX_LEN - 1 + + +def _encode_dspark_exact_expected_yields( + *, exact_local, num_cells: int, + yield_scale: int) -> Tuple[int, Tuple[int, ...]]: + """Encode one rank's exact yields without splitting the collective.""" + if yield_scale <= 0: + raise ValueError("DSpark exact-policy yield scale must be positive") + num_cells = int(num_cells) + if num_cells < 0: + raise ValueError("DSpark exact-policy cell count must be nonnegative") + if exact_local is None: + return 0, tuple(-1 for _ in range(num_cells)) + + native = float(exact_local.native_expected_yield) + native_valid = math.isfinite(native) and native >= 0.0 + native_wire = int(math.ceil(native * yield_scale)) if native_valid else 0 + local_values = tuple(exact_local.compact_expected_yields) + compact_wires = [] + for cell_index in range(num_cells): + value = (float(local_values[cell_index]) + if cell_index < len(local_values) else math.nan) + if (not native_valid or not math.isfinite(value) or value < 0.0 + or value > native): + compact_wires.append(-1) + else: + compact_wires.append(int(math.floor(value * yield_scale))) + return native_wire, tuple(compact_wires) + + +def _decode_dspark_exact_expected_yields( + *, + payloads: List[List[int]], + exact_cells: Iterable[Tuple[int, int]], + graph_batch_size: int, + yield_scale: int, +) -> Tuple[float, Dict[int, float], Dict[int, float]]: + """Decode aggregate yields and worst-rank loss per real request.""" + if yield_scale <= 0: + raise ValueError("DSpark exact-policy yield scale must be positive") + native_expected_yield = sum( + float(payload[_DSPARK_EXACT_NATIVE_YIELD_INDEX]) + for payload in payloads) / yield_scale + compact_expected_yields = {} + compact_max_yield_losses_per_request = {} + for cell_index, (cell_graph_batch_size, + verifier_budget) in enumerate(exact_cells): + if cell_graph_batch_size != graph_batch_size: + continue + rank_yields = [ + float(payload[_DSPARK_EXACT_WIRE_PREFIX_LEN + cell_index]) + for payload in payloads + ] + infeasible = any(value < 0.0 for value in rank_yields) + compact_expected_yields[verifier_budget] = (0.0 if infeasible else + sum(rank_yields) / + yield_scale) + active_rank_losses = [] + if not infeasible: + for payload, compact_yield_wire in zip(payloads, rank_yields): + num_real_requests = int(payload[1]) + if num_real_requests <= 0: + continue + native_yield_wire = float( + payload[_DSPARK_EXACT_NATIVE_YIELD_INDEX]) + if compact_yield_wire > native_yield_wire: + raise RuntimeError( + "DSpark exact-policy compact yield exceeds native yield" + ) + active_rank_losses.append( + (native_yield_wire - compact_yield_wire) / yield_scale / + num_real_requests) + compact_max_yield_losses_per_request[verifier_budget] = max( + active_rank_losses, default=0.0) + return ( + native_expected_yield, + compact_expected_yields, + compact_max_yield_losses_per_request, + ) + + +def _validate_dspark_exact_bucket(*, exact_shape, bucket) -> None: + """Require the graph key to match the globally selected exact cell.""" + if exact_shape is None: + return + expected_bucket = int(exact_shape[1]) + if bucket is None or int(bucket) != expected_bucket: + raise RuntimeError( + "DSpark exact policy did not publish its selected captured " + f"verifier bucket V={expected_bucket}; got {bucket}") + + +def _validate_dspark_adp_acceptance_gate( + *, confidence_enabled: bool, attention_dp_enabled: bool, + speculation_gate_enabled: bool) -> None: + """Reject a rank-local gate that can split the ADP collective protocol.""" + if confidence_enabled and attention_dp_enabled and speculation_gate_enabled: + raise ValueError( + "DSpark confidence scheduling with attention DP does not support " + "acceptance_rate_window_size/acceptance_rate_threshold because " + "the per-rank gate can break collective ordering") + + +def _classify_dspark_exact_generation_rows( + generation_requests: List[LlmRequest], +) -> Tuple[Optional[List[LlmRequest]], bool]: + """Return logical exact rows and whether one idle ADP dummy is scheduled.""" + if not generation_requests: + return None, False + real_requests = [ + request for request in generation_requests if not request.is_dummy + ] + if len(real_requests) == len(generation_requests): + return real_requests, False + one_adp_dummy = bool(not real_requests and len(generation_requests) == 1 + and generation_requests[0].is_attention_dp_dummy) + if one_adp_dummy: + return real_requests, True + return None, False + + +def _dspark_exact_common_graph_batch_size( + payloads: List[List[int]], round_up: Callable[[int], int]) -> int: + """Resolve a common exact G, preserving zero as all-rank termination.""" + widest_rows = max((int(payload[1]) for payload in payloads), default=0) + return 0 if widest_rows <= 0 else int(round_up(widest_rows)) + + +def _dspark_exact_secondary_padding_ready( + payloads: List[List[int]], + *, + graph_batch_size: int, + verifier_budget: int, + secondary_ready_index: int, +) -> bool: + """Whether every zero-real peer can represent the selected remainder.""" + if verifier_budget % graph_batch_size <= 1: + return True + return all( + int(payload[1]) != 0 or bool(int(payload[secondary_ready_index])) + for payload in payloads) class _ADPForwardIntent(IntEnum): @@ -639,6 +784,12 @@ def __init__( threshold = getattr(spec_config, 'acceptance_rate_threshold', None) if window and threshold is not None: self.speculation_gate = SpeculationGate(window, threshold) + _validate_dspark_adp_acceptance_gate( + confidence_enabled=getattr(self.model_engine, + "_dspark_confidence_enabled", False), + attention_dp_enabled=self.enable_attention_dp, + speculation_gate_enabled=self.speculation_gate is not None, + ) # response used data self.response_lock = threading.Lock() @@ -2090,7 +2241,7 @@ def get_latest_iteration_stats(self): Returns the per-iterations statistics computed since last call to this method. Contains at most iter_stats_max_iterations iterations. """ - if self.enable_iter_perf_stats == False: + if not self.enable_iter_perf_stats: return [] latest_stats = (IterationStats(), None) @@ -3931,13 +4082,414 @@ def _drain_relay_sends_before_forward(self, microbatch_id: int) -> None: self.dist.pp_size) % self.num_micro_batches self.wait_on_pp_send_handles(self.send_handles, due_microbatch_id) + def _stage_dspark_confidence(self, worker, planner, gen_requests, + mixed) -> None: + buffer = worker.staged_confidence_buffer() + if buffer is None: + return + worker.bump_draft_seq() + if mixed is not None and mixed.costs is not None: + from .dspark_mixed_runtime import stage_mixed_ownership + + planner.stage_confidence(buffer, worker.confidence_stamp_buffer()) + stage_mixed_ownership(planner, worker, int(self.iter_counter), + gen_requests) + else: + planner.stage_confidence(buffer) + + def _dspark_confidence_draft_len(self, + scheduled_batch: ScheduledRequests) -> int: + """Agree one DSpark verify shape and stage confidence asynchronously.""" + max_draft_len = self.model_engine.max_draft_len + worker = self.model_engine._get_spec_worker() + planner = getattr(worker, "verify_planner", None) + if planner is None: + return max_draft_len + mixed = self.model_engine._dspark_mixed_runtime + if mixed is not None: + mixed.reset_step() + + runtime_cost_table = getattr(self.model_engine, + "_dspark_sps_cost_table", None) + if runtime_cost_table is not None and planner.cost_table is not runtime_cost_table: + planner.install_runtime_cost_table(runtime_cost_table) + + gen_requests = scheduled_batch.generation_requests + is_distributed = self.dist is not None and self.dist.tp_size > 1 + compute_windows = bool( + getattr(self.model_engine, "_dspark_trims_submitted_tokens", False)) + + ragged_lens = None + device_budget = None + exact_local = None + exact_table = planner.exact_cost_table + if exact_table is None: + raise RuntimeError( + "DSpark confidence scheduling has no authenticated exact SPS table" + ) + exact_cells = tuple( + getattr(self.model_engine, "_dspark_exact_candidate_cells", ())) + if not exact_cells: + exact_cells = exact_table.candidate_cells() + exact_requests: Optional[List[LlmRequest]] = gen_requests + exact_zero_real_adp_dummy = False + exact_requests, exact_zero_real_adp_dummy = ( + _classify_dspark_exact_generation_rows(gen_requests)) + + # A stale capacity must never re-rank a later batch. + self.model_engine._dspark_device_budget = None + if compute_windows: + if scheduled_batch.context_requests: + pass + else: + if exact_requests is not None: + exact_rows = [ + worker.confidence_row_for(request.py_request_id) + for request in exact_requests + ] + exact_local = planner.prepare_exact_sps_decision( + num_gen_requests=len(exact_requests), rows=exact_rows) + + local_rows = (exact_local.num_requests if exact_local is not None else + len(ragged_lens) if ragged_lens is not None else 0) + local_tokens = (sum( + 1 + int(value) + for value in ragged_lens) if ragged_lens is not None else 0) + local_payload = [ + 1 if exact_local is not None or ragged_lens is not None else 0, + int(local_rows), + int(local_tokens), + 1 if scheduled_batch.can_run_cuda_graph else 0, + ] + + exact_yield_scale = 1_000_000 + exact_identity_words = (0, ) * 8 + exact_identity_words = tuple( + getattr( + self.model_engine, + "_dspark_exact_identity_words", + exact_table.collective_identity_words, + )) + native_yield_wire, compact_yields_wire = _encode_dspark_exact_expected_yields( + exact_local=exact_local, + num_cells=len(exact_cells), + yield_scale=exact_yield_scale, + ) + local_payload.extend([ + 1, + 1 if exact_local is not None else 0, + len(exact_cells), + *exact_identity_words, + native_yield_wire, + *compact_yields_wire, + ]) + + runner = self.model_engine.cuda_graph_runner + local_payload.extend([ + int(scheduled_batch.batch_size), + 1 if planner.max_verify_len in getattr( + runner, "padding_dummy_requests", {}) else 0, + 1 if planner.max_verify_len in getattr( + runner, "secondary_padding_dummy_requests", {}) else 0, + ]) + if mixed is not None: + local_payload.extend( + mixed.local_vote( + scheduled_batch, + worker, + planner, + int(self.iter_counter), + int(self.dist.tp_rank) if is_distributed else 0, + executor_eligible=bool( + compute_windows + and self._dspark_adp_shape_cache_enabled()), + )) + payloads = (self.dist.tp_allgather_int64(local_payload).tolist() + if is_distributed else [local_payload]) + + exact_payload_end = _DSPARK_EXACT_WIRE_PREFIX_LEN + len(exact_cells) + expected_payload_len = exact_payload_end + _DSPARK_ADP_WIRE_TRAILER_LEN + mixed_peers = None + if mixed is not None: + payloads, mixed_peers = mixed.split_votes(payloads, + expected_payload_len) + if any(len(payload) != expected_payload_len for payload in payloads): + raise RuntimeError( + "DSpark ADP agreement payload has an inconsistent shape") + + exact_table_flags = [int(payload[4]) for payload in payloads] + if len(set(exact_table_flags)) != 1: + raise RuntimeError( + "DSpark ranks disagree on whether an exact SPS table is installed" + ) + if exact_table_flags[0]: + exact_cell_counts = [int(payload[6]) for payload in payloads] + exact_identities = [ + tuple(int(value) for value in payload[7:15]) + for payload in payloads + ] + if (len(set(exact_cell_counts)) != 1 + or exact_cell_counts[0] != len(exact_cells) + or len(set(exact_identities)) != 1): + raise RuntimeError( + "DSpark exact SPS ranks disagree on the authenticated " + "(G,V), T(G,V), or minimum-gain table identity") + + all_can_graph = all(int(payload[3]) for payload in payloads) + peer_stats = None + if not all(int(payload[0]) for payload in payloads): + ragged_lens = None + exact_local = None + else: + peer_stats = [[ + int(payload[1]), + int(payload[2]), 1 if all_can_graph else 0 + ] for payload in payloads] + + batch_size_index = exact_payload_end + padding_dummy_index = exact_payload_end + 1 + secondary_padding_dummy_index = exact_payload_end + 2 + peer_batch_sizes = tuple( + int(payload[batch_size_index]) for payload in payloads) + widest_batch_size = max(peer_batch_sizes, default=0) + common_padded_batch_size = runner._round_up_batch_size_with_draft_len( + widest_batch_size, planner.max_verify_len) + padding_needed = any( + int(payload[batch_size_index]) != common_padded_batch_size + for payload in payloads) + padding_supported = bool( + not padding_needed + or (runner.padding_enabled + and common_padded_batch_size <= runner.config.batch_size)) + padding_ready = bool( + common_padded_batch_size > 0 and padding_supported and all( + int(payload[batch_size_index]) == common_padded_batch_size + or int(payload[padding_dummy_index]) for payload in payloads)) + exact_policy_agreed = bool(all(int(payload[5]) for payload in payloads)) + agreement = ADPShapeAgreement( + iteration=int(self.iter_counter), + batch_identity=id(scheduled_batch), + local_batch_size=int(scheduled_batch.batch_size), + peer_batch_sizes=peer_batch_sizes, + all_can_graph=all_can_graph, + widest_batch_size=widest_batch_size, + graph_batch_size=common_padded_batch_size, + draft_len=int(planner.max_verify_len), + padding_ready=padding_ready, + ) + runner.adp_shape_agreement = agreement + if (mixed_peers is not None and agreement.can_queue + and mixed.apply(scheduled_batch, worker, planner, mixed_peers, + int(self.iter_counter))): + # Common context-aware get_padding_params supplies the BREAKABLE + # body. A decode-only peer must not replay its local full graph. + runner.adp_shape_agreement = dataclasses.replace( + agreement, + all_can_graph=False, + ragged_bucket=None, + reuse_graph_shape=False, + ) + self._stage_dspark_confidence(worker, planner, gen_requests, mixed) + return int(planner.max_verify_len) + if not agreement.can_queue: + runner.agreed_ragged_bucket = None + runner.ragged_pad_verify_len = 0 + runner.ragged_zero_real_high_rows = 0 + for request in gen_requests: + request.py_verify_len = None + return int(planner.max_verify_len) + + exact_shape = None + if exact_local is not None: + if not all_can_graph: + exact_local = None + else: + if any(not int(payload[5]) for payload in payloads): + raise RuntimeError( + "DSpark exact SPS collective payloads disagree on the " + "measured (G,V) grid") + common_graph_batch_size = _dspark_exact_common_graph_batch_size( + payloads, runner._round_up_batch_size) + if common_graph_batch_size == 0: + exact_local = None + else: + from ..speculative.dspark_planner import \ + select_exact_sps_candidate + + ( + native_expected_yield, + compact_expected_yields, + compact_max_yield_losses_per_request, + ) = _decode_dspark_exact_expected_yields( + payloads=payloads, + exact_cells=exact_cells, + graph_batch_size=common_graph_batch_size, + yield_scale=exact_yield_scale, + ) + selected_verifier_budget = select_exact_sps_candidate( + graph_batch_size=common_graph_batch_size, + native_expected_yield=native_expected_yield, + compact_expected_yields=compact_expected_yields, + compact_max_yield_losses_per_request=( + compact_max_yield_losses_per_request), + cost_table=exact_table, + ) + if selected_verifier_budget == 0 or not ( + _dspark_exact_secondary_padding_ready( + payloads, + graph_batch_size=common_graph_batch_size, + verifier_budget=selected_verifier_budget, + secondary_ready_index= + secondary_padding_dummy_index, + )): + exact_local = None + else: + ragged_lens, local_exact_budget, exact_pad_tokens = ( + planner.allocate_exact_sps_candidate( + exact_local, + graph_batch_size=common_graph_batch_size, + verifier_budget=selected_verifier_budget, + )) + if exact_local.num_requests == 0: + if not exact_zero_real_adp_dummy: + raise RuntimeError( + "zero-real exact SPS allocation has no " + "authenticated attention-DP dummy row") + _, remainder = divmod(selected_verifier_budget, + common_graph_batch_size) + scheduled_tokens = int(exact_pad_tokens) + int( + remainder > 0) + ragged_lens = [scheduled_tokens - 1] + exact_shape = ( + common_graph_batch_size, + selected_verifier_budget, + exact_pad_tokens, + ) + if planner.device_windows and local_exact_budget > 0: + device_budget = int(local_exact_budget) + + from ..speculative.dspark_verify import \ + _exact_cell_geometry + + peer_stats = [] + for payload in payloads: + peer_rows = int(payload[1]) + peer_geometry = _exact_cell_geometry( + num_real=peer_rows, + graph_batch_size=common_graph_batch_size, + verifier_budget=selected_verifier_budget, + min_verify_len=planner.cfg.min_verify_len, + max_verify_len=planner.max_verify_len, + ) + if peer_geometry is None: + raise RuntimeError( + "Selected exact SPS cell is infeasible for " + f"peer rows={peer_rows}, " + f"G={common_graph_batch_size}, " + f"V={selected_verifier_budget}") + _, peer_real_tokens, _ = peer_geometry + peer_stats.append([ + peer_rows, peer_real_tokens, + 1 if all_can_graph else 0 + ]) + + bucket = None + ragged_active = False + if ragged_lens is not None: + if exact_shape is None: + raise RuntimeError( + "DSpark exact allocation produced no measured graph shape") + if all_can_graph: + bucket = self.model_engine.fit_ragged_verify_lens( + gen_requests, + ragged_lens, + peer_stats=peer_stats, + exact_shape=exact_shape, + exact_zero_real=exact_zero_real_adp_dummy, + ) + ragged_active = bucket is not None + if ragged_active and device_budget is not None: + self.model_engine._dspark_device_budget = int(device_budget) + else: + runner.agreed_ragged_bucket = None + runner.ragged_pad_verify_len = 0 + runner.ragged_zero_real_high_rows = 0 + for request, window in zip(gen_requests, ragged_lens): + request.py_verify_len = int(window) + ragged_active = True + + reuse_graph_shape = bool(exact_policy_agreed and runner.enabled + and all_can_graph and padding_ready + and not runner.config.use_mrope) + _validate_dspark_exact_bucket(exact_shape=exact_shape, bucket=bucket) + if (reuse_graph_shape and exact_shape is not None + and exact_shape[0] != common_padded_batch_size): + raise RuntimeError( + "DSpark exact policy selected a graph batch size that " + "disagrees with the cached ADP padding shape") + runner.adp_shape_agreement = dataclasses.replace( + agreement, + ragged_bucket=None if bucket is None else int(bucket), + reuse_graph_shape=reuse_graph_shape, + ) + + if not ragged_active: + for request in gen_requests: + request.py_verify_len = None + runner.agreed_ragged_bucket = None + runner.ragged_pad_verify_len = 0 + runner.ragged_zero_real_high_rows = 0 + + self._stage_dspark_confidence(worker, planner, gen_requests, mixed) + return int(planner.max_verify_len) + def _handle_dynamic_draft_len(self, scheduled_batch: ScheduledRequests) -> None: - """Synchronize draft length and buffers before preparing resources.""" - update_draft_len(self.model_engine, - scheduled_batch, - speculation_permanently_disabled=self. - speculation_permanently_disabled) + """Synchronize draft width and confidence policy before KV admission.""" + spec_config = getattr(self.model_engine, "spec_config", None) + confidence_driven = bool( + spec_config is not None + and getattr(spec_config, "enable_confidence_scheduling", False)) + if not confidence_driven or self.speculation_permanently_disabled: + # Main owns the static, scheduled and disabled draft-buffer rules. + update_draft_len(self.model_engine, + scheduled_batch, + speculation_permanently_disabled=self. + speculation_permanently_disabled) + return + + signature = ( + int(self.iter_counter), + id(scheduled_batch), + int(scheduled_batch.batch_size), + int(scheduled_batch.num_context_requests), + int(scheduled_batch.num_generation_requests), + ) + if getattr(self, "_dspark_dynamic_handled_signature", + None) == signature: + return + runtime_draft_len = self._dspark_confidence_draft_len(scheduled_batch) + agreement = self.model_engine.cuda_graph_runner.adp_shape_agreement + if (agreement is not None and agreement.matches( + scheduled_batch, self.iter_counter, padded=False) + and not agreement.can_queue): + self.model_engine.runtime_draft_len = runtime_draft_len + else: + update_draft_len(self.model_engine, + scheduled_batch, + draft_len=runtime_draft_len) + self._dspark_dynamic_handled_signature = signature + + def _dspark_adp_shape_cache_enabled(self) -> bool: + runner = getattr(self.model_engine, "cuda_graph_runner", None) + return bool( + self.enable_attention_dp and self.dist is not None + and self.dist.tp_size > 1 and runner is not None + and getattr(self.model_engine, "_dspark_confidence_enabled", False) + and getattr(self, "kv_connector_manager", None) is None + and getattr(self, "kv_cache_transceiver", None) is None + and getattr(self, "drafter", None) is None + and getattr(self, "_pp_rebalance_drain_iters", None) is None) def _initialize_remote_tail_adp(self) -> None: """Reserve the execution-only idle row after final KV/graph warmup.""" @@ -4013,6 +4565,28 @@ def _can_queue(self, scheduled_batch): # can_queue_this_rank is for case that the batch is not empty on this rank, but empty on other ranks # For bs == 1, we cannot pad dummy request to make the batch non-empty since it will cause the batch size to be 2. # 1 for dummy request, 1 for the yet-to-complete but not-yet-updated request. + runner = getattr(self.model_engine, "cuda_graph_runner", None) + if self._dspark_adp_shape_cache_enabled(): + signature = ( + int(self.iter_counter), + id(scheduled_batch), + int(scheduled_batch.batch_size), + int(scheduled_batch.num_context_requests), + int(scheduled_batch.num_generation_requests), + ) + if getattr(self, "_dspark_dynamic_handled_signature", + None) != signature: + runner.adp_shape_agreement = None + self._handle_dynamic_draft_len(scheduled_batch) + agreement = runner.adp_shape_agreement + if agreement is not None and agreement.matches( + scheduled_batch, self.iter_counter, padded=False): + return agreement.can_queue, agreement.can_queue_this_rank + elif runner is not None: + # Optional surfaces can mutate the request set after policy + # selection, so never reuse their preceding agreement. + runner.adp_shape_agreement = None + if getattr(self, "_enable_remote_tail_adp", False): # An idle execution row needs no active-request or sampler slot. # Still veto an all-empty fleet (e.g. cancellations after the vote). @@ -5612,7 +6186,7 @@ def _maybe_finish_pp_rebalance(self) -> None: self._pp_rebalance_drain_iters = None def _suspend_padding_dummies_for_rebalance( - self, mgr: KVCacheManagerV2) -> List[Tuple[int, LlmRequest]]: + self, mgr: KVCacheManagerV2) -> List[Tuple[int, int, LlmRequest]]: """Suspend the CUDA-graph padding dummies before ``adjust()``. ``CUDAGraphRunner`` retains one padding dummy per captured draft @@ -5629,21 +6203,27 @@ def _suspend_padding_dummies_for_rebalance( and silently drop padded batches to eager mode for the rest of the process lifetime. - Returns the ``(draft_len, dummy)`` pairs that were suspended. + Returns ``(draft_len, variant, dummy)`` triples. Variant zero is the + low-row dummy and variant one is the high-row dummy. """ runner = getattr(self.model_engine, "cuda_graph_runner", None) if runner is None: return [] - suspended: List[Tuple[int, LlmRequest]] = [] - for draft_len, dummy in runner.padding_dummy_requests.items(): - if mgr.is_request_active(dummy.py_request_id): - mgr.suspend_request(dummy) - suspended.append((draft_len, dummy)) + suspended: List[Tuple[int, int, LlmRequest]] = [] + caches = ( + getattr(runner, "padding_dummy_requests", {}), + getattr(runner, "secondary_padding_dummy_requests", {}), + ) + for variant, cache in enumerate(caches): + for draft_len, dummy in cache.items(): + if mgr.is_request_active(dummy.py_request_id): + mgr.suspend_request(dummy) + suspended.append((draft_len, variant, dummy)) return suspended def _resume_padding_dummies_after_rebalance( self, mgr: KVCacheManagerV2, - suspended: List[Tuple[int, LlmRequest]]) -> None: + suspended: List[Tuple[int, int, LlmRequest]]) -> None: """Resume the padding dummies suspended for ``adjust()``. A real request that fails to resume is left suspended for the @@ -5661,7 +6241,8 @@ def _resume_padding_dummies_after_rebalance( runner = getattr(self.model_engine, "cuda_graph_runner", None) if runner is None: return - for draft_len, dummy in suspended: + failed_draft_lens = set() + for draft_len, variant, dummy in suspended: if mgr.resume_request(dummy): continue if getattr(self, "_enable_remote_tail_adp", @@ -5671,8 +6252,13 @@ def _resume_padding_dummies_after_rebalance( continue logger.warning( "Could not resume the CUDA graph padding dummy request " - f"(draft_len={draft_len}) after a KV pool rebalance; " + f"(draft_len={draft_len}, variant={variant}) after a KV pool " + "rebalance; " "releasing it and falling back to lazy re-creation.") + failed_draft_lens.add(draft_len) + # Resume both siblings before releasing a low/high pair. Releasing + # inline could free a resumed sibling and later try to resume it. + for draft_len in sorted(failed_draft_lens): runner.release_padding_dummy(self.resource_manager, draft_len) @contextmanager @@ -6356,6 +6942,10 @@ def _configure_csa2_remote_tail(self, request: LlmRequest) -> None: request.py_csa2_remote_tail_split = split def _validate_request(self, request: LlmRequest): + if getattr(self.model_engine, "_dspark_confidence_enabled", False): + from ..speculative.dspark import validate_confidence_request + + validate_confidence_request(self.model_engine.spec_config, request) self._configure_csa2_remote_tail(request) if (isinstance(self.kv_cache_manager, KVCacheManagerV2) and request.kv_cache_retention_config is not None): @@ -9048,7 +9638,7 @@ def _enqueue_responses(self, responses: Iterable[Tuple[int, LlmResponse]]): # (TODO: joyang) There are other types of responses, we need to sort out. if type( resp - ) == LlmResponse and req_id in self.result_wait_queues and self.result_wait_queues[ + ) is LlmResponse and req_id in self.result_wait_queues and self.result_wait_queues[ req_id] is not None: self.result_wait_queues[req_id].put_response.remote( resp.client_id, resp) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 348328d75d96..83a3e3cbd00b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -11,7 +11,7 @@ from contextlib import contextmanager from dataclasses import dataclass from itertools import chain -from typing import Optional, Tuple +from typing import Callable, Optional, Tuple import torch from strenum import StrEnum @@ -69,6 +69,19 @@ FP4_MLA_TOKENS_PER_BLOCK = 128 +@dataclass(frozen=True) +class _RuntimeConstructors: + """Explicit internal constructors for a non-serving offline executor. + + Normal serving passes no bundle and retains the ordinary constructors. + This is not an LLM argument and cannot enable confidence admission. + """ + + model_engine: Callable[..., PyTorchModelEngine] + executor: type[PyExecutor] + kv_cache_creator: type[KvCacheCreator] + + class _ExecutorMemoryMonitor: """Currently this focuses on tracking memory usage and related errors.""" @@ -484,6 +497,7 @@ def _create_py_executor_impl( tokenizer: Optional[TokenizerBase] = None, profiling_stage_data: Optional[dict] = None, resource_governor_queue=None, + runtime_constructors: Optional[_RuntimeConstructors] = None, ) -> PyExecutor: """Create and initialize a PyExecutor instance from the given LLM arguments. @@ -515,6 +529,7 @@ def _create_py_executor_impl( tokenizer=tokenizer, profiling_stage_data=profiling_stage_data, resource_governor_queue=resource_governor_queue, + runtime_constructors=runtime_constructors, ) py_executor.metrics.update(creation_metrics) return py_executor @@ -528,6 +543,7 @@ def _create_py_executor( tokenizer: Optional[TokenizerBase] = None, profiling_stage_data: Optional[dict] = None, resource_governor_queue=None, + runtime_constructors: Optional[_RuntimeConstructors] = None, ) -> PyExecutor: """Create and initialize a PyExecutor while recording scoped metrics.""" initial_model_engine_metrics: dict[str, float | dict[str, float]] = {} @@ -790,7 +806,9 @@ def allocation_scope(current_stage: ExecutorMemoryType): model_weights_restore_mode = sleep_config.restore_modes[ ExecutorMemoryType.MODEL_WEIGHTS_MAIN] - model_engine = PyTorchModelEngine( + model_engine_type = (PyTorchModelEngine if runtime_constructors is None + else runtime_constructors.model_engine) + model_engine = model_engine_type( model_path=checkpoint_dir, llm_args=llm_args, mapping=mapping, @@ -1130,7 +1148,9 @@ def allocation_scope(current_stage: ExecutorMemoryType): # Get draft config for one-engine speculative decoding if available draft_config = getattr(model_engine.model, 'draft_config', None) - kv_cache_creator = KvCacheCreator( + kv_cache_creator_type = (KvCacheCreator if runtime_constructors is None + else runtime_constructors.kv_cache_creator) + kv_cache_creator = kv_cache_creator_type( model_engine=model_engine, draft_model_engine=draft_model_engine, mapping=mapping, @@ -1196,6 +1216,8 @@ def allocation_scope(current_stage: ExecutorMemoryType): # run gc.collect() to free memory of the previous py_executor, avoid cudaFree overlap with cuda graph capture gc.collect() py_executor = create_py_executor_instance( + executor_type=None + if runtime_constructors is None else runtime_constructors.executor, dist=dist, resources=resources, mapping=mapping, @@ -1277,6 +1299,8 @@ def allocation_scope(current_stage: ExecutorMemoryType): # run gc.collect() to free memory of the previous py_executor, avoid cudaFree overlap with cuda graph capture gc.collect() py_executor = create_py_executor_instance( + executor_type=None if runtime_constructors is None else + runtime_constructors.executor, dist=dist, resources=resources, mapping=mapping, @@ -1341,6 +1365,7 @@ def create_py_executor( tokenizer: Optional[TokenizerBase] = None, profiling_stage_data: Optional[dict] = None, resource_governor_queue=None, + runtime_constructors: Optional[_RuntimeConstructors] = None, ) -> PyExecutor: """Create a PyExecutor and roll back a partially initialized DWDP runtime.""" previous_dwdp_manager = get_global_dwdp_manager() @@ -1354,6 +1379,7 @@ def create_py_executor( tokenizer=tokenizer, profiling_stage_data=profiling_stage_data, resource_governor_queue=resource_governor_queue, + runtime_constructors=runtime_constructors, ) except BaseException: current_dwdp_manager = get_global_dwdp_manager() diff --git a/tensorrt_llm/_torch/speculative/dspark.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_cost_profile.py b/tensorrt_llm/_torch/speculative/dspark_cost_profile.py new file mode 100644 index 000000000000..db8a9ded3e96 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_cost_profile.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Explicit non-serving DSpark graph-cell profiling contracts.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Literal + +from pydantic import ConfigDict, Field, StrictBool, StrictInt, model_validator + +from tensorrt_llm.llmapi.utils import StrictBaseModel + +__all__ = ["DSparkCostProfileConfig", "DSparkProfileCell"] + + +class DSparkProfileCell(StrictBaseModel): + """One rank-local graph shape; zero budget denotes the native full block.""" + + model_config = ConfigDict(frozen=True) + graph_batch_size: StrictInt = Field(gt=0, description="Rank-local request count G.") + verifier_budget: StrictInt = Field( + ge=0, description="Verifier token count V; zero means native." + ) + + +class DSparkCostProfileConfig(StrictBaseModel): + """Construction-only graph grid, never a serving admission configuration.""" + + model_config = ConfigDict(frozen=True) + physical_k: StrictInt = Field(default=3, gt=0, description="Full physical draft block K.") + cells: tuple[DSparkProfileCell, ...] = Field(min_length=1, description="Declared graph cells.") + record_confidence_head: StrictBool = Field( + default=True, description="Execute the head in graphs." + ) + measurement_scope: Literal["graph-replay", "whole-step"] = Field( + default="graph-replay", description="Original replay or executor-loop cost boundary." + ) + max_samples: StrictInt = Field( + default=4096, gt=0, le=65536, description="Maximum raw replay records for the run." + ) + + @model_validator(mode="after") + def _validate_grid(self) -> DSparkCostProfileConfig: + seen = set() + native = set() + for cell in self.cells: + pair = (cell.graph_batch_size, cell.verifier_budget) + if pair in seen: + raise ValueError("Profile cells must not contain duplicates") + seen.add(pair) + if cell.verifier_budget == 0: + native.add(cell.graph_batch_size) + elif ( + self.physical_k == 1 + or not cell.graph_batch_size + <= cell.verifier_budget + < cell.graph_batch_size * (self.physical_k + 1) + ): + raise ValueError("Compact cells require K>1 and G <= V < G*(K+1); native uses V=0") + if native != {cell.graph_batch_size for cell in self.cells}: + raise ValueError("Every profiled G requires a V=0 native control") + return self + + def validate_graph_ladder(self, graph_batch_sizes: Sequence[int]) -> None: + """Require the exact runtime ladder, not an undocumented subset.""" + if any(type(value) is not int or value <= 0 for value in graph_batch_sizes): + raise ValueError("Runtime graph row counts must be positive integers") + if len(set(graph_batch_sizes)) != len(graph_batch_sizes): + raise ValueError("Runtime graph ladder contains duplicates") + if set(graph_batch_sizes) != {cell.graph_batch_size for cell in self.cells}: + raise ValueError("Profile grid must cover the exact captured graph ladder") + + def graph_shapes(self, graph_batch_sizes: Sequence[int]) -> list[tuple[int, ...]]: + """Return native and positive-V capture shapes in deterministic order.""" + self.validate_graph_ladder(graph_batch_sizes) + return [ + (cell.graph_batch_size, self.physical_k) + if cell.verifier_budget == 0 + else (cell.graph_batch_size, self.physical_k, cell.verifier_budget) + for cell in sorted( + self.cells, key=lambda item: (item.graph_batch_size, item.verifier_budget) + ) + ] + + def require_cell(self, cell: DSparkProfileCell) -> None: + """Reject an undeclared shape before request metadata can change.""" + if cell not in self.cells: + raise ValueError("Requested profile cell is not in the construction grid") diff --git a/tensorrt_llm/_torch/speculative/dspark_cost_profiler.py b/tensorrt_llm/_torch/speculative/dspark_cost_profiler.py new file mode 100644 index 000000000000..1b3da52abfc4 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_cost_profiler.py @@ -0,0 +1,563 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSpark chain adapter and explicit offline driver; no serving admission.""" + +from __future__ import annotations + +import hashlib +import importlib +import importlib.metadata +import math +import sys +import time +from functools import partial +from pathlib import Path +from typing import Callable + +import torch +import yaml +from mpi4py import MPI + +from tensorrt_llm._torch.pyexecutor.dspark_request_profile import create_profile_executor +from tensorrt_llm._torch.speculative.dspark_cost_profile import DSparkCostProfileConfig +from tensorrt_llm._utils import mpi_comm +from tensorrt_llm.executor.request import GenerationRequest +from tensorrt_llm.executor.worker import GenerationExecutorWorker +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig, TorchLlmArgs +from tensorrt_llm.sampling_params import SamplingParams +from tensorrt_llm.tools.dspark_cost_export import ( + observed_engine_fingerprint, + source_origin_from_args, + source_origin_input_metadata, +) +from tensorrt_llm.tools.verify_costs import ( + canonical_sha256, + create_private_output_directory, + isolated_control_communicator, + parse_json, + sha256_file, + write_new_json, +) + +__all__ = ["run_profile"] + + +class OwnedInputWorker(GenerationExecutorWorker): + """Construct the owned factory inside each original worker, not its proxy.""" + + def __init__( + self, + *, + profile: DSparkCostProfileConfig, + **kwargs, + ) -> None: + self._profile = profile + if kwargs.get("rpc_addr") is not None: + raise ValueError("Offline profiling cannot expose a serving RPC endpoint") + super().__init__(**kwargs) + + def _get_pytorch_executor_factory(self) -> Callable: + # Retain every original worker setup operation and its ordering. + return partial(create_profile_executor, profile=self._profile) + + +def response_rows(worker, inputs, output_length, timeout): + """Use original request conversion/enqueue/futures, with real native IDs.""" + deadline = time.monotonic() + timeout + futures = [] + for row in inputs: + if len(row["prompt_token_ids"]) + output_length > worker.max_seq_len: + raise ValueError("Original max_seq_len would shorten the fixed output length") + sampling = SamplingParams( + max_tokens=output_length, + temperature=0.0, + top_k=1, + n=1, + best_of=1, + ignore_eos=True, + detokenize=False, + exclude_input_from_output=True, + ) + request = GenerationRequest( + prompt_token_ids=list(row["prompt_token_ids"]), + sampling_params=sampling, + streaming=False, + ) + if time.monotonic() >= deadline: + raise TimeoutError("Original total child request budget exhausted before enqueue") + future = worker.submit(request) + native_id = worker._client_id_to_request_id[request.id] + futures.append((row, request.id, native_id, future)) + results = [] + for row, client_id, native_id, future in futures: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("Original total child request budget exhausted before response") + result = future.result(timeout=remaining) + outputs = result.outputs + if ( + len(outputs) != 1 + or outputs[0].finish_reason != "length" + or len(outputs[0].token_ids) != output_length + or any(type(token) is not int or token < 0 for token in outputs[0].token_ids) + ): + raise RuntimeError("Original response did not complete the fixed greedy output") + results.append( + { + "input_id": row["input_id"], + "client_id": client_id, + "native_request_id": native_id, + "prompt_token_ids": row["prompt_token_ids"], + "output_token_ids": list(outputs[0].token_ids), + "finish_reason": outputs[0].finish_reason, + } + ) + return results + + +def _collective(control, label, callback): + error, result = None, None + try: + result = callback() + except Exception as exception: + error = type(exception).__name__ + ": " + str(exception) + errors = control.allgather(error) + if any(error is not None for error in errors): + raise RuntimeError(label + ": " + repr(errors)) + return result + + +def _identity(worker, args, profile, expected_files): + """Record live installed sources/runtime and adapter-defined execution layout.""" + modules = ( + "tensorrt_llm._torch.pyexecutor.model_engine", + "tensorrt_llm._torch.pyexecutor._util", + "tensorrt_llm._torch.pyexecutor.py_executor_creator", + "tensorrt_llm._torch.pyexecutor.cuda_graph_runner", + "tensorrt_llm._torch.pyexecutor.py_executor", + "tensorrt_llm.executor.base_worker", + "tensorrt_llm.executor.worker", + "tensorrt_llm.llmapi.llm_args", + "tensorrt_llm._torch.speculative.dspark", + "tensorrt_llm._torch.pyexecutor.dspark_profile_engine", + "tensorrt_llm._torch.pyexecutor.dspark_request_profile", + "tensorrt_llm._torch.speculative.dspark_cost_profile", + "tensorrt_llm._torch.speculative.dspark_cost_profiler", + "tensorrt_llm.tools.verify_costs", + ) + if profile.measurement_scope == "whole-step": + modules += ("tensorrt_llm.tools.dspark_cost_export",) + sources = {} + for name in modules: + module = importlib.import_module(name) + sources[name] = sha256_file(Path(module.__file__)) + native_paths = set(torch.ops.loaded_libraries) + for name, module in tuple(sys.modules.items()): + if name.startswith(("tensorrt_llm.", "torch._C")): + path = getattr(module, "__file__", None) + if path is not None and Path(path).suffix in (".so", ".pyd"): + native_paths.add(path) + native = {str(Path(path).resolve()): sha256_file(Path(path)) for path in sorted(native_paths)} + if not native: + raise RuntimeError("Loaded native library identity is missing") + package = importlib.metadata.metadata("tensorrt_llm") + source_urls = [ + url for url in package.get_all("Project-URL", []) if url.startswith("Source Commit,") + ] + device = torch.cuda.get_device_properties(torch.cuda.current_device()) + versions = { + "torch": str(torch.__version__), + "cuda": torch.version.cuda, + "tensorrt_llm": importlib.metadata.version("tensorrt_llm"), + } + actual_files = { + "inputs_sha256": sha256_file(args.inputs), + "model_config_sha256": sha256_file(args.model / "config.json"), + } + if profile.measurement_scope == "whole-step": + actual_files["recorded_source_origin"] = source_origin_from_args(args) + actual_files["source_origin_observation"] = source_origin_input_metadata( + getattr(args, "recorded_source_origin", None), + getattr(args, "recorded_source_origin_sha256", None), + ) + if actual_files != expected_files: + raise ValueError("Profile inputs or model configuration changed after validation") + mapping = worker.engine.model_engine.mapping + model = worker.engine.model_engine.model + attention_classes = sorted( + { + type(module).__module__ + "." + type(module).__qualname__ + for module in model.modules() + if type(module).__module__.startswith("tensorrt_llm._torch.attention") + } + ) + if profile.measurement_scope == "whole-step": + # Retain observed model/module overlays as well as the adapter roster. + # Native library bytes are already recorded independently above. + for instance in model.modules(): + name = type(instance).__module__ + if name.startswith("tensorrt_llm.") and name not in sources: + module = importlib.import_module(name) + path = getattr(module, "__file__", None) + if path is not None and Path(path).suffix == ".py": + sources[name] = sha256_file(Path(path)) + return { + "identity_scope": "DIAGNOSTIC_PARTIAL_LOADED_RUNTIME", + "production_live_fingerprint_qualified": False, + "loaded_source_sha256": sources, + "package_source_commit_urls": source_urls, + "runtime": versions, + "local_native_libraries": native, + "local_device": { + "name": device.name, + "major": device.major, + "minor": device.minor, + "total_memory": device.total_memory, + }, + "topology": { + "tp": mapping.tp_size, + "pp": mapping.pp_size, + "cp": mapping.cp_size, + "ep": mapping.moe_ep_size, + "attention_dp": mapping.enable_attention_dp, + }, + "runtime_configuration": { + "max_batch_size": worker.llm_args.max_batch_size, + "max_seq_len": worker.max_seq_len, + "max_num_tokens": worker.llm_args.max_num_tokens, + "disable_overlap_scheduler": worker.llm_args.disable_overlap_scheduler, + "enable_iter_perf_stats": worker.llm_args.enable_iter_perf_stats, + "iter_perf_stats_interval": worker.llm_args.iter_perf_stats_interval, + "cuda_graph_config": worker.llm_args.cuda_graph_config.model_dump(mode="json"), + "kv_cache_config": worker.llm_args.kv_cache_config.model_dump(mode="json"), + "configured_attention_backend": str(worker.llm_args.attn_backend), + }, + **actual_files, + "profile": profile.model_dump(mode="json"), + "generations_per_shape": args.generations, + "sampling": { + "output_tokens": args.output_tokens, + "temperature": 0.0, + "top_k": 1, + "ignore_eos": True, + "beam_width": 1, + }, + "declared_shape_ids": [canonical_sha256(_shape(cell, profile)) for cell in profile.cells], + "observed_attention_module_classes": attention_classes, + "kernel_dispatch_qualification": False, + } + + +def _shape(cell, profile) -> dict: + """DSpark-specific shape identity; never reduce different attention layouts.""" + return { + "workload": "dspark_chain", + "physical_k": profile.physical_k, + "graph_batch_size": cell.graph_batch_size, + "verifier_budget": cell.verifier_budget, + "attention_layout": "uniform" if cell.verifier_budget == 0 else "packed_chain_rows", + "captured_head": profile.record_confidence_head, + } + + +def _raw_generation(generation, cell, repetition, identity, rank, world): + """The DSpark adapter decodes lifecycle records; the shared reducer does not.""" + samples = {row["sample_index"]: row for row in generation["observation"]["samples"]} + output = [] + for join in generation["lifecycle_joins"]: + index = join["sample_index"] + if index is None: + continue + sample = samples.pop(index) + scope = identity["profile"].get("measurement_scope", "graph-replay") + duration = ( + join["host_step_time_ms"] + if scope == "whole-step" + else sample["replay_entry_elapsed_ms"] + ) + output.append( + { + "sample_index": index, + "executor_iteration": join["executor_iteration"], + "elapsed_ms": duration, + "graph_replay_diagnostic_ms": sample["replay_entry_elapsed_ms"], + "eligible": all(join["sampler_update_applied"]), + "request_ids": sample["request_ids"], + "graph_key": sample["graph_key"], + } + ) + if samples: + raise RuntimeError("Unjoined graph replay samples") + graph_keys = {canonical_sha256(row["graph_key"]) for row in output} + if len(graph_keys) != 1: + raise RuntimeError("Prescribed cell observed zero or incompatible graph replay keys") + profile = DSparkCostProfileConfig(**identity["profile"]) + shape = _shape(cell, profile) + return { + "schema": "verify-cost-raw-v1", + "rank": rank, + "world_size": world, + "generation": repetition, + "generation_closed": True, + "model_shutdown_completed": False, + "identity": identity, + "timing_scope": ( + "original-executor-loop-hostStepTimeMS" + if profile.measurement_scope == "whole-step" + else "original-graph-replay-entry-including-static-copies" + ), + "declared_shape_id": canonical_sha256(shape), + "execution_shape": {**shape, "observed_graph_key": output[0]["graph_key"]}, + "samples": output, + "lifecycle": generation, + "policy_overhead_measured": False, + "production_admission_qualified": False, + } + + +def _prepare(args, world): + """Validate rank-local inputs before any peer enters collective construction.""" + if world > 1 and MPI.Query_thread() < MPI.THREAD_MULTIPLE: + raise ValueError("Concurrent runtime/profile MPI requires THREAD_MULTIPLE") + if ( + args.generations < 1 + or args.output_tokens < 1 + or not math.isfinite(args.timeout) + or args.timeout <= 0 + ): + raise ValueError("Positive generations, output length and total timeout required") + grid = parse_json(args.grid.read_bytes()) + if not isinstance(grid, dict): + raise ValueError("The profile grid must be a JSON object") + if "physical_k" in grid and grid["physical_k"] != args.physical_k: + raise ValueError("CLI physical K and grid differ") + cost_scope = getattr(args, "cost_scope", "graph-replay") + if "measurement_scope" in grid and grid["measurement_scope"] != cost_scope: + raise ValueError("CLI measurement scope and grid differ") + profile = DSparkCostProfileConfig( + **{**grid, "physical_k": args.physical_k, "measurement_scope": cost_scope} + ) + config = yaml.safe_load(args.config.read_text(encoding="utf-8")) + if not isinstance(config, dict): + raise ValueError("The runtime configuration must be a YAML mapping") + llm_args = TorchLlmArgs(**{**config, "model": str(args.model), "backend": "pytorch"}) + if profile.measurement_scope == "whole-step" and ( + not llm_args.enable_iter_perf_stats or getattr(args, "source_root", None) is None + ): + raise ValueError("Whole-step collection requires enable_iter_perf_stats and --source-root") + spec = llm_args.speculative_config + mapping = llm_args.parallel_config.to_mapping() + if profile.measurement_scope == "whole-step" and world > 1 and not mapping.enable_attention_dp: + raise ValueError( + "Schema-v2 whole-step export currently requires attention-DP fanout or one rank" + ) + if ( + not isinstance(spec, DSparkDecodingConfig) + or spec.max_draft_len != profile.physical_k + or mapping.tp_size != world + or mapping.pp_size != 1 + or mapping.cp_size != 1 + ): + raise ValueError("Actual MPI topology/model physical K must match the profile") + if ( + llm_args.guided_decoding_backend is not None + or llm_args.batched_logits_processor is not None + or llm_args.lora_config is not None + or llm_args.cache_transceiver_config is not None + or llm_args.kv_connector_config is not None + or llm_args.enable_in_graph_sampling + or llm_args.encode_only + or llm_args.enable_resource_governor + ): + raise ValueError("The offline driver requires aggregate greedy generation without adapters") + input_bytes = args.inputs.read_bytes() + inputs = [parse_json(line) for line in input_bytes.splitlines() if line.strip()] + factor = world if mapping.enable_attention_dp else 1 + maximum = max(cell.graph_batch_size for cell in profile.cells) * factor + if len(inputs) < maximum or any( + not isinstance(row, dict) + or not isinstance(row.get("prompt_token_ids"), list) + or not row["prompt_token_ids"] + or any(type(token) is not int or token < 0 for token in row["prompt_token_ids"]) + for row in inputs + ): + raise ValueError("Enough nonempty original integer-token prompts are required") + for index, row in enumerate(inputs): + row["input_id"] = index + files = { + "inputs_sha256": hashlib.sha256(input_bytes).hexdigest(), + "model_config_sha256": sha256_file(args.model / "config.json"), + } + if profile.measurement_scope == "whole-step": + files["recorded_source_origin"] = source_origin_from_args(args) + files["source_origin_observation"] = source_origin_input_metadata( + getattr(args, "recorded_source_origin", None), + getattr(args, "recorded_source_origin_sha256", None), + ) + fingerprint = canonical_sha256( + { + "config": config, + "profile": profile.model_dump(mode="json"), + "inputs": inputs, + "generations": args.generations, + "output_tokens": args.output_tokens, + "timeout": args.timeout, + "output_directory": str(args.output.resolve()), + **files, + } + ) + return profile, llm_args, inputs, factor, fingerprint, files + + +def run_profile(args) -> None: + """Use original worker requests and state updates for a prescribed offline grid.""" + # Original request-count/payload broadcasts stay on mpi_comm(). Main-thread + # profile collectives need a separate context before any worker can start. + with isolated_control_communicator(mpi_comm()) as control: + _run_profile(args, control) + + +def _run_profile(args, control) -> None: + """Keep the isolated control context alive through startup and all closure.""" + rank, world = control.Get_rank(), control.Get_size() + profile, llm_args, inputs, factor, fingerprint, expected_files = _collective( + control, "prepare", lambda: _prepare(args, world) + ) + fingerprints = control.allgather(fingerprint) + if len(set(fingerprints)) != 1: + raise ValueError("Ranks disagree on inputs/configuration/grid") + _collective( + control, + "output directory", + lambda: create_private_output_directory(args.output) if rank == 0 else None, + ) + worker = None + records = [] + succeeded = False + shutdown_completed = False + try: + + def construct(): + nonlocal worker + # Preserve a successful local owner even if a peer reports failure + # at the following collective error-agreement boundary. + worker = OwnedInputWorker( + profile=profile, + engine=None, + hf_model_dir=args.model, + tokenizer=None, + llm_args=llm_args, + is_llm_executor=True, + ) + return worker + + _collective(control, "construct", construct) + identity = _collective( + control, "identity", lambda: _identity(worker, args, profile, expected_files) + ) + identity["devices"] = control.allgather(identity.pop("local_device")) + identity["rank_native_libraries"] = control.allgather( + identity.pop("local_native_libraries") + ) + if profile.measurement_scope == "whole-step": + identity["recorded_engine_fingerprint"] = _collective( + control, + "recorded cost engine metadata", + lambda: observed_engine_fingerprint( + identity, + source_root=args.source_root, + world_size=world, + recorded_origin=getattr(args, "recorded_source_origin", None), + recorded_origin_sha256=getattr(args, "recorded_source_origin_sha256", None), + ), + ) + if len(set(control.allgather(canonical_sha256(identity)))) != 1: + raise ValueError("Actual rank identities differ") + deadline = time.monotonic() + args.timeout + generation_id = 0 + for cell in profile.cells: + selected = inputs[: factor * cell.graph_batch_size] + for repetition in range(args.generations): + _collective( + control, + "begin", + lambda: worker.engine.begin_profile_generation(generation_id, cell), + ) + responses, error = None, None + if rank == 0: + try: + responses = response_rows( + worker, + selected, + args.output_tokens, + max(0, deadline - time.monotonic()), + ) + except Exception as exception: + error = type(exception).__name__ + ": " + str(exception) + error, responses = control.bcast((error, responses), root=0) + if error is not None: + raise RuntimeError(error) + generation = _collective( + control, + "drain", + lambda: worker.engine.end_profile_generation(control=control), + ) + + def join_generation(): + by_id = {row["native_request_id"]: row for row in responses} + if len(by_id) != len(responses): + raise RuntimeError("Duplicate native response IDs") + for final in generation["original_final_requests"]: + response = by_id[final["request_id"]] + if final["tokens"] != ( + response["prompt_token_ids"] + response["output_token_ids"] + ): + raise RuntimeError("Final native request tokens and response differ") + return _raw_generation(generation, cell, repetition, identity, rank, world) + + record = _collective(control, "generation join", join_generation) + record["responses"] = responses + records.append(record) + generation_id += 1 + succeeded = True + finally: + original_failure = sys.exc_info()[1] + try: + # Every rank enters the same control phase, including a rank whose + # constructor failed. Do not manufacture a model-shutdown receipt + # for a worker that was never successfully constructed. + _collective( + control, "shutdown", lambda: worker.shutdown() if worker is not None else None + ) + shutdown_completed = worker is not None + finally: + for index, record in enumerate(records): + record["model_shutdown_completed"] = shutdown_completed + record["run_complete"] = succeeded and shutdown_completed + _collective( + control, + "preserve raw", + lambda: write_new_json( + args.output / f"rank-{rank}-record-{index}.json", record + ), + ) + if not succeeded or not shutdown_completed: + _collective( + control, + "preserve failure", + lambda: write_new_json( + args.output / f"rank-{rank}-failure.json", + { + "schema": "verify-cost-run-failure-v1", + "run_complete": False, + "failure_type": ( + type(original_failure).__name__ + if original_failure + else "ShutdownFailure" + ), + "model_shutdown_completed": shutdown_completed, + }, + ), + ) + if not succeeded: + raise RuntimeError("Offline generation did not complete") diff --git a/tensorrt_llm/_torch/speculative/dspark_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..e9c6222ba8fe 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,42 +377,59 @@ 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. Only + configured graph tiers are installed as runtime execution choices; unused + measured tiers remain in the original returned payload. """ payload = _read_sps_cost_payload(path) schema_version = payload.get("schema_version") 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, + runtime_graph_batch_sizes=graph_batch_sizes, + ), payload def _build_exact_sps_cost_table( payload: dict[str, object], *, max_draft_len: int, + validate_artifact_freshness: bool = True, + runtime_graph_batch_sizes: Sequence[int] | None = None, ) -> ExactSpsCostTable: exact_tables = payload["cost_tables"] assert isinstance(exact_tables, dict) parsed_tables: dict[int, ExactSpsCostRow] = {} + # Scope execution choices, not serialized measurements. The loader already + # validates the entire original artifact, including unused finite cells. + runtime_graphs = set(runtime_graph_batch_sizes) if runtime_graph_batch_sizes is not None else None for graph_batch_size, exact_payload in exact_tables.items(): canonical_graph_batch_size = _parse_graph_batch_size_key(graph_batch_size) + if runtime_graphs is not None and canonical_graph_batch_size not in runtime_graphs: + continue assert isinstance(exact_payload, dict) parsed_tables[canonical_graph_batch_size] = ExactSpsCostRow( token_counts=tuple( @@ -452,6 +475,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 +722,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 +736,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. @@ -720,7 +746,8 @@ def validate_sps_cost_table_payload( Raises: TypeError: If the artifact has malformed container types. ValueError: If its schema, fingerprint, coverage, or provenance does - not exactly describe the configured graph ladder. + not cover the configured graph ladder. Strict freshness additionally + requires the original full graph ladder and exact live identity. """ if not isinstance(payload, dict): raise TypeError("SPS cost artifact must contain a JSON object") @@ -733,30 +760,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) @@ -778,19 +812,22 @@ def validate_sps_cost_table_payload( field="iteration drain tail_graph_batch_size", minimum=1, ) - if tail_graph_batch_size not in measured_tables: + if (tail_graph_batch_size not in measured_tables + or tail_graph_batch_size not in configured_graph_batch_sizes): raise ValueError( - "Iteration drain guard requires a measured native tail table for " + "Iteration drain guard requires a captured, measured native tail table for " f"G={tail_graph_batch_size}" ) fingerprint_graph_batch_sizes = set(fingerprint["rank_local_graph_batch_sizes"]) if ( - configured_graph_batch_sizes != set(measured_tables) - or configured_graph_batch_sizes != fingerprint_graph_batch_sizes + set(measured_tables) != fingerprint_graph_batch_sizes + or not configured_graph_batch_sizes.issubset(measured_tables) + or (validate_freshness and configured_graph_batch_sizes != set(measured_tables)) ): raise ValueError( - "Multi-G SPS graph batch sizes must match exactly across runtime " - "graphs, cost_tables, and engine_fingerprint" + "Runtime graphs require measured coverage; cost_tables and their " + "original engine_fingerprint must match exactly. Strict freshness " + "also requires the original runtime graph ladder." ) runtime_max_draft_len = _require_json_int( max_draft_len, field="runtime max_draft_len", minimum=1 @@ -876,7 +913,7 @@ def validate_sps_cost_table_payload( gpu_count = fingerprint["gpu_count"] global_graph_batch_sizes = set(fingerprint["global_graph_batch_sizes"]) if global_graph_batch_sizes != { - graph_batch_size * gpu_count for graph_batch_size in configured_graph_batch_sizes + graph_batch_size * gpu_count for graph_batch_size in measured_tables }: raise ValueError( "SPS engine_fingerprint global_graph_batch_sizes are inconsistent " @@ -1130,8 +1167,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 +1236,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_reproduction.py b/tensorrt_llm/_torch/speculative/dspark_reproduction.py new file mode 100644 index 000000000000..109339584460 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_reproduction.py @@ -0,0 +1,370 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSpark adapter for explicit offline collection and ordinary automatic policy.""" + +from __future__ import annotations + +import copy +import importlib +import importlib.metadata +import math +import os +import sys +import time +from pathlib import Path + +import torch +import yaml +from mpi4py import MPI + +from tensorrt_llm._torch.pyexecutor.dspark_observation_runtime import observation_factory +from tensorrt_llm._torch.speculative.dspark_cost_profiler import _collective, response_rows +from tensorrt_llm._utils import mpi_comm +from tensorrt_llm.executor.worker import GenerationExecutorWorker +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig, TorchLlmArgs +from tensorrt_llm.tools.dspark_cost_export import ( + source_origin_from_args, + source_origin_input_metadata, +) +from tensorrt_llm.tools.dspark_observations import validate_responses +from tensorrt_llm.tools.verify_costs import ( + canonical_sha256, + create_private_output_directory, + isolated_control_communicator, + parse_json, + sha256_file, + write_new_json, +) + +_POLICY_FIELDS = ( + "enable_confidence_scheduling", + "enable_fused_confidence_scheduler", + "validate_confidence_artifact_freshness", + "confidence_sts_path", + "confidence_sps_table_path", + "confidence_sps_live_fingerprint_path", + "confidence_admission_receipt_path", + "confidence_admission_receipt_sha256", +) + + +class ObservationWorker(GenerationExecutorWorker): + def __init__(self, *, collect_calibration=False, **kwargs): + self._collect_calibration = collect_calibration + if kwargs.get("rpc_addr") is not None: + raise ValueError("Offline collection/benchmark cannot expose serving RPC") + super().__init__(**kwargs) + + def _get_pytorch_executor_factory(self): + return observation_factory(collect_calibration=self._collect_calibration) + + +def _prepare(args, world): + if os.environ.get("TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS"): + raise ValueError("Collection and benchmarking require real target acceptance") + if ( + args.physical_k < 1 + or args.batch_size < 1 + or args.output_tokens < 1 + or args.warmup_cohorts < 0 + or not math.isfinite(args.timeout) + or args.timeout <= 0 + ): + raise ValueError("Positive K/BS/output/full budget and nonnegative warmups required") + if world > 1 and MPI.Query_thread() < MPI.THREAD_MULTIPLE: + raise ValueError("Concurrent ordinary worker/control MPI requires THREAD_MULTIPLE") + config = yaml.safe_load(args.config.read_text(encoding="utf-8")) + if not isinstance(config, dict) or not isinstance(config.get("speculative_config"), dict): + raise ValueError("Explicit public DSpark speculative_config required") + spec = config["speculative_config"] + if spec.get("max_draft_len") != args.physical_k: + raise ValueError("CLI physical K differs from actual configuration") + numerical = copy.deepcopy(config) + for field in _POLICY_FIELDS: + numerical["speculative_config"].pop(field, None) + if args.mode == "dynamic": + if args.physical_k == 1 or args.cost_table is None: + raise ValueError("Dynamic chain needs K>1 and a structurally valid measured cost table") + spec["enable_confidence_scheduling"] = True + spec["confidence_sps_table_path"] = str(args.cost_table.resolve(strict=True)) + if args.calibration is not None: + spec["confidence_sts_path"] = str(args.calibration.resolve(strict=True)) + else: + for field in _POLICY_FIELDS: + spec.pop(field, None) + spec["enable_confidence_scheduling"] = False + if args.cost_table is not None or args.calibration is not None: + raise ValueError("Static/calibration collection does not consume dynamic artifacts") + llm_args = TorchLlmArgs(**{**config, "model": str(args.model), "backend": "pytorch"}) + mapping = llm_args.parallel_config.to_mapping() + factor = world if mapping.enable_attention_dp else 1 + if ( + not isinstance(llm_args.speculative_config, DSparkDecodingConfig) + or not llm_args.speculative_config.draft_is_embedded_in_target + or mapping.tp_size != world + or mapping.pp_size != 1 + or mapping.cp_size != 1 + or args.batch_size % factor + or llm_args.max_batch_size < args.batch_size // factor + or llm_args.cuda_graph_config is None + or not llm_args.cuda_graph_config.enable_padding + ): + raise ValueError( + "K/topology/BS capacity and padded decode graphs must match the request cohort" + ) + if ( + llm_args.guided_decoding_backend is not None + or llm_args.batched_logits_processor is not None + or llm_args.lora_config is not None + or llm_args.cache_transceiver_config is not None + or llm_args.kv_connector_config is not None + or llm_args.enable_in_graph_sampling + or llm_args.encode_only + or llm_args.enable_resource_governor + ): + raise ValueError( + "This fixed-cohort driver supports ordinary aggregate greedy generation only" + ) + if args.mode == "collect" and llm_args.disable_overlap_scheduler: + raise ValueError( + "This source-paired collector currently uses the original overlap device-input seam" + ) + data = args.inputs.read_bytes() + inputs = [parse_json(line) for line in data.splitlines() if line.strip()][: args.batch_size] + if len(inputs) != args.batch_size or any( + not isinstance(row, dict) + or not isinstance(row.get("prompt_token_ids"), list) + or not row["prompt_token_ids"] + or any(type(token) is not int or token < 0 for token in row["prompt_token_ids"]) + for row in inputs + ): + raise ValueError("One nonempty integer-token prompt per selected request required") + for index, row in enumerate(inputs): + row["input_id"] = index + if len(row["prompt_token_ids"]) + args.output_tokens > llm_args.max_seq_len: + raise ValueError("Fixed output would exceed original max_seq_len") + origin = source_origin_from_args(args) + contract = { + "physical_k": args.physical_k, + "batch_size": args.batch_size, + "output_tokens": args.output_tokens, + "inputs_sha256": canonical_sha256(inputs), + "model_config_sha256": sha256_file(args.model / "config.json"), + "numerical_configuration_sha256": canonical_sha256(numerical), + "warmup_cohorts": args.warmup_cohorts, + "sampling": {"temperature": 0.0, "top_k": 1, "ignore_eos": True, "beam_width": 1}, + } + return llm_args, inputs, contract, origin + + +def _observed_identity(worker, args, origin): + # Finite source/native/runtime provenance. HEAD is only recorded base; + # observed module bytes include this overlay and are never relabeled old. + modules = ( + "tensorrt_llm._torch.speculative.dspark", + "tensorrt_llm._torch.speculative.dspark_verify", + "tensorrt_llm._torch.speculative.dspark_reproduction", + "tensorrt_llm._torch.pyexecutor.model_engine", + "tensorrt_llm._torch.pyexecutor.py_executor", + "tensorrt_llm._torch.pyexecutor.cuda_graph_runner", + "tensorrt_llm._torch.pyexecutor.dspark_observation_runtime", + "tensorrt_llm._torch.speculative.spec_sampler_base", + "tensorrt_llm.tools.dspark_observations", + "tensorrt_llm.llmapi.llm_args", + ) + loaded = {name: sha256_file(Path(importlib.import_module(name).__file__)) for name in modules} + if source_origin_from_args(args) != origin: + raise ValueError("Observed source changed during model construction") + for name, digest in loaded.items(): + path = args.source_root / (name.replace(".", "/") + ".py") + if sha256_file(path) != digest: + raise ValueError("Actual loaded Python bytes differ from recorded source origin") + native_paths = set(torch.ops.loaded_libraries) + for name, module in tuple(sys.modules.items()): + path = getattr(module, "__file__", None) + if ( + name.startswith(("tensorrt_llm.", "torch._C")) + and path + and Path(path).suffix in (".so", ".pyd") + ): + native_paths.add(path) + native = {str(Path(path).resolve()): sha256_file(Path(path)) for path in sorted(native_paths)} + if not native: + raise ValueError("Actual native runtime identity missing") + device = torch.cuda.get_device_properties(torch.cuda.current_device()) + mapping = worker.engine.model_engine.mapping + return { + **origin, + "source_origin_observation": source_origin_input_metadata( + getattr(args, "recorded_source_origin", None), + getattr(args, "recorded_source_origin_sha256", None), + ), + "loaded_source_sha256": loaded, + "runtime": { + "torch": str(torch.__version__), + "cuda": torch.version.cuda, + "tensorrt_llm": importlib.metadata.version("tensorrt_llm"), + }, + "local_native_libraries": native, + "local_device": {"name": device.name, "major": device.major, "minor": device.minor}, + "topology": { + "tp": mapping.tp_size, + "ep": mapping.moe_ep_size, + "attention_dp": mapping.enable_attention_dp, + "pp": mapping.pp_size, + }, + "artifact_sha256": { + name: sha256_file(path) + for name, path in (("cost", args.cost_table), ("calibration", args.calibration)) + if path + }, + } + + +def run_native(args): + with isolated_control_communicator(mpi_comm()) as control: + return _run_native(args, control) + + +def _run_native(args, control): + rank, world = control.Get_rank(), control.Get_size() + deadline = time.monotonic() + args.timeout + llm_args, inputs, contract, origin = _collective( + control, "prepare", lambda: _prepare(args, world) + ) + if len(set(control.allgather(canonical_sha256((contract, origin))))) != 1: + raise ValueError("Ranks disagree on request/configuration/source inputs") + _collective( + control, + "fresh outputs", + lambda: create_private_output_directory(args.output) if rank == 0 else None, + ) + worker, record, payload = None, None, None + succeeded, shutdown_completed = False, False + try: + + def construct(): + nonlocal worker + worker = ObservationWorker( + collect_calibration=args.mode == "collect", + engine=None, + hf_model_dir=args.model, + tokenizer=None, + llm_args=llm_args, + is_llm_executor=True, + ) + return worker + + _collective(control, "construct", construct) + identity = _collective( + control, "observed identity", lambda: _observed_identity(worker, args, origin) + ) + identity["devices"] = control.allgather(identity.pop("local_device")) + identity["rank_native_libraries"] = control.allgather( + identity.pop("local_native_libraries") + ) + # Only measured runtime/source fields determine matched identity; + # artifact origin metadata is preserved separately, not rewritten. + contract["observed_runtime_sha256"] = canonical_sha256( + {key: value for key, value in identity.items() if key != "artifact_sha256"} + ) + for cohort in range(args.warmup_cohorts + 1): + _collective( + control, + "reset prior cohort routing", + lambda: worker.engine.observation_receipt(reset=True, control=control), + ) + control.Barrier() + started = time.monotonic() + responses, error = None, None + if rank == 0: + try: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("Full original startup/model budget exhausted") + responses = response_rows(worker, inputs, args.output_tokens, remaining) + except Exception as exception: + error = type(exception).__name__ + ": " + str(exception) + error, responses, elapsed = control.bcast( + (error, responses, time.monotonic() - started), root=0 + ) + if error: + raise RuntimeError(error) + _collective( + control, + "responses", + lambda: validate_responses(responses, inputs, args.output_tokens), + ) + receipt = _collective( + control, + "original drain", + lambda: worker.engine.observation_receipt(control=control), + ) + if cohort == args.warmup_cohorts: + routing = control.allgather(receipt["routing"]) + combined = { + name: sum(row[name] for row in routing) + for name in ("native_replays", "compact_replays") + } + record = { + "schema": "dspark-public-run-v1", + "mode": args.mode, + "calibration_collection": args.mode == "collect", + "inputs": inputs, + "responses": responses, + "elapsed_seconds": elapsed, + "comparison_contract": contract, + "identity": identity, + "resolved_confidence_enabled": receipt["resolved_confidence_enabled"], + "runtime_confidence_enabled": receipt["runtime_confidence_enabled"], + "routing": combined, + "rank_graph_keys": control.allgather(receipt["graph_keys"]), + "task_accuracy_qualified": False, + "production_admission_qualified": False, + } + if args.mode == "collect": + record["rank_drain"] = control.allgather(receipt["drain"]) + if args.mode == "collect": + payload = _collective( + control, "paired full-K payload", worker.engine._native_pairs.payload + ) + payload["meta"]["identity"] = identity + succeeded = True + finally: + failure = sys.exc_info()[1] + try: + _collective(control, "original shutdown", lambda: worker.shutdown() if worker else None) + shutdown_completed = worker is not None + finally: + if record is not None: + record["run_complete"] = succeeded and shutdown_completed + record["model_shutdown_completed"] = shutdown_completed + _collective( + control, + "preserve record", + lambda: write_new_json(args.output / f"rank-{rank}-run.json", record), + ) + if succeeded and shutdown_completed and payload is not None: + payload["meta"]["model_shutdown_completed"] = True + + def preserve_pairs(): + with (args.output / f"rank-{rank}-pairs.pt").open("xb") as handle: + torch.save(payload, handle) + + _collective(control, "preserve pairs", preserve_pairs) + if not succeeded or not shutdown_completed: + _collective( + control, + "preserve failure", + lambda: write_new_json( + args.output / f"rank-{rank}-failure.json", + { + "schema": "dspark-public-failure-v1", + "run_complete": False, + "failure_type": type(failure).__name__ + if failure + else "ShutdownFailure", + "model_shutdown_completed": shutdown_completed, + }, + ), + ) diff --git a/tensorrt_llm/_torch/speculative/dspark_schedule.py b/tensorrt_llm/_torch/speculative/dspark_schedule.py index e9bd692b6259..2e49e64791b6 100644 --- a/tensorrt_llm/_torch/speculative/dspark_schedule.py +++ b/tensorrt_llm/_torch/speculative/dspark_schedule.py @@ -27,12 +27,16 @@ from dataclasses import dataclass import torch +import triton +import triton.language as tl __all__ = [ "DSparkScheduleConfig", + "DSparkFusedScheduleError", "HOST_POLICY_WINDOWS_SNAPSHOT_OUTPUT", "NATIVE_UNIFORM_VERIFY_OUTPUT", "compute_survival", + "schedule_verify_lens_topk_fused_fill", "schedule_verify_lens_topk", ] @@ -50,6 +54,15 @@ NATIVE_UNIFORM_VERIFY_OUTPUT = "native_uniform_verify" +class DSparkFusedScheduleError(RuntimeError): + """A synchronous optional fused-scheduler dispatch failure. + + Input and graph-contract errors remain ordinary ``ValueError`` exceptions; + callers may safely disable one fused G only for this wrapped launch error + and rerun the established tensor implementation. + """ + + @dataclass(frozen=True) class DSparkScheduleConfig: """Bounds for the per-request verify length. @@ -179,3 +192,269 @@ def schedule_verify_lens_topk( granted = selected.view(bs, schedulable).sum(dim=1).to(torch.int32) verify_lens += granted return torch.clamp(verify_lens, min=floor, max=cfg.resolved_max_verify_len) + + +@triton.jit(do_not_specialize=["num_real", "budget"]) +def _schedule_topk_rank_kernel( + survival_ptr, + output_ptr, + stride_row, + stride_position, + num_real, + budget, + survival_eps, + FLOOR: tl.constexpr, + COLS: tl.constexpr, + NUM_CANDIDATES: tl.constexpr, + BLOCK_COLS: tl.constexpr, + BLOCK_CP: tl.constexpr, +): + request = tl.program_id(0) + position = tl.arange(0, BLOCK_COLS) + candidate_mask = position < COLS + survival = tl.load( + survival_ptr + request * stride_row + (position + FLOOR).to(tl.int64) * stride_position, + mask=candidate_mask, + other=0.0, + ).to(tl.float32) + valid = candidate_mask & (survival >= survival_eps) + masked_survival = tl.where(valid, survival, float("-inf")) + + rank = tl.zeros([BLOCK_COLS], dtype=tl.int32) + for peer_start in range(0, NUM_CANDIDATES, BLOCK_CP): + peer = peer_start + tl.arange(0, BLOCK_CP) + peer_mask = peer < NUM_CANDIDATES + peer_request = peer // COLS + peer_position = peer % COLS + peer_survival = tl.load( + survival_ptr + + peer_request.to(tl.int64) * stride_row + + (peer_position + FLOOR).to(tl.int64) * stride_position, + mask=peer_mask, + other=0.0, + ).to(tl.float32) + peer_valid = peer_mask & (peer_request < num_real) & (peer_survival >= survival_eps) + peer_masked_survival = tl.where(peer_valid, peer_survival, float("-inf")) + + higher = peer_masked_survival[None, :] > masked_survival[:, None] + equal = peer_masked_survival[None, :] == masked_survival[:, None] + earlier = (peer_position[None, :] < position[:, None]) | ( + (peer_position[None, :] == position[:, None]) & (peer_request[None, :] < request) + ) + before = (higher | (equal & earlier)) & peer_valid[None, :] + rank += tl.sum(before.to(tl.int32), axis=1) + + selected = valid & (rank < budget) + tl.store(output_ptr + request, tl.sum(selected.to(tl.int32), axis=0)) + + +@triton.jit(do_not_specialize=["num_real", "budget", "pad_len", "graph_num_tokens"]) +def _schedule_topk_fill_kernel( + output_ptr, + num_rows, + num_real, + budget, + pad_len, + graph_num_tokens, + TOKEN_FLOOR: tl.constexpr, + MAX_TOKEN_LEN: tl.constexpr, + HAS_CANDIDATES: tl.constexpr, + BLOCK: tl.constexpr, +): + rows = tl.arange(0, BLOCK) + row_mask = rows < num_rows + is_real = row_mask & (rows < num_real) + selected_extra = tl.load( + output_ptr + rows, + mask=is_real & (budget > 0) & HAS_CANDIDATES, + other=0, + ).to(tl.int32) + token_len = tl.where(is_real, TOKEN_FLOOR + selected_extra, pad_len) + + # Reproduce fill_bucket_device's exact-pad real-row phase. The explicit + # policy budget controls only the confidence-ranked allocation. Any + # already-paid graph remainder is filled round-robin in request-index + # order; fusion must not promote it into another confidence budget. + spare = graph_num_tokens - tl.sum(tl.where(row_mask, token_len, 0), axis=0) + for _ in range(MAX_TOKEN_LEN - 1): + headroom = is_real & (token_len < MAX_TOKEN_LEN) + in_cycle = tl.cumsum(headroom.to(tl.int32), axis=0) <= spare + grant = headroom & in_cycle + token_len += grant.to(tl.int32) + spare -= tl.sum(grant.to(tl.int32), axis=0) + tl.store(output_ptr + rows, token_len, mask=row_mask) + + +def _schedule_verify_lens_topk_fused_fill_triton( + *, + survival: torch.Tensor, + budget: int, + num_real: int, + pad_len: int, + graph_num_tokens: int, + cfg: DSparkScheduleConfig, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Triton implementation of policy top-k plus exact round-robin fill.""" + num_rows = survival.shape[0] + if num_rows == 0: + return torch.empty(0, dtype=torch.int32, device=survival.device) if out is None else out[:0] + if num_rows > 256: + raise ValueError(f"fused top-k scheduler supports at most 256 rows, got {num_rows}") + schedulable = int(cfg.schedulable_per_request) + output = ( + torch.empty(num_rows, dtype=torch.int32, device=survival.device) + if out is None + else out[:num_rows] + ) + + # Specialize rank work by the captured graph row count, never by the + # step-varying real-row count. This keeps one compiled kernel per all-G + # graph key; peer rows past num_real are masked inside the kernel. + num_candidates = num_rows * schedulable + ranked = budget > 0 and num_real > 0 and num_candidates > 0 + if ranked: + block_cols = triton.next_power_of_2(schedulable) + block_cp = 128 + try: + _schedule_topk_rank_kernel[(num_real,)]( + survival, + output, + survival.stride(0), + survival.stride(1), + num_real, + budget, + float(cfg.survival_eps), + FLOOR=int(cfg.min_verify_len), + COLS=schedulable, + NUM_CANDIDATES=num_candidates, + BLOCK_COLS=block_cols, + BLOCK_CP=block_cp, + ) + except Exception as exc: + raise DSparkFusedScheduleError("DSpark fused top-k rank launch failed") from exc + finalize_block = triton.next_power_of_2(num_rows) + try: + _schedule_topk_fill_kernel[(1,)]( + output, + num_rows, + num_real, + budget, + pad_len, + graph_num_tokens, + TOKEN_FLOOR=int(cfg.min_verify_len) + 1, + MAX_TOKEN_LEN=int(cfg.resolved_max_verify_len) + 1, + HAS_CANDIDATES=schedulable > 0, + BLOCK=finalize_block, + ) + except Exception as exc: + raise DSparkFusedScheduleError("DSpark fused round-robin fill launch failed") from exc + return output + + +def schedule_verify_lens_topk_fused_fill( + *, + survival: torch.Tensor, + budget: int, + num_real: int, + pad_len: int, + cfg: DSparkScheduleConfig, + graph_num_tokens: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Fuse the established top-k schedule and exact round-robin bucket fill. + + The explicit ``budget`` retains exactly the semantics of + :func:`schedule_verify_lens_topk`. Pad rows receive exactly ``pad_len`` + tokens; any graph-token remainder is granted to real rows in index-order + round-robin cycles, exactly as ``fill_bucket_device(..., pad_fill=...)``. + CUDA uses two deterministic Triton launches. CPU executes the established + tensor operations as the functional oracle. When ``out`` is supplied, the + first ``num_rows`` int32 elements are overwritten and returned; CUDA then + performs no output allocation. + """ + if survival.dim() != 2: + raise ValueError(f"survival must be [bs, K], got shape {tuple(survival.shape)}") + num_rows, num_positions = survival.shape + if out is not None: + if out.device != survival.device: + raise ValueError("out must share survival.device") + if out.dtype != torch.int32: + raise TypeError(f"out must have dtype torch.int32, got {out.dtype}") + if out.dim() != 1 or out.numel() < num_rows: + raise ValueError(f"out must be a 1-D tensor with at least {num_rows} elements") + if not out.is_contiguous(): + raise ValueError("out must be contiguous") + if num_positions != cfg.block_size: + raise ValueError( + f"survival has {num_positions} positions but block_size is {cfg.block_size}" + ) + budget = int(budget) + num_real = int(num_real) + pad_len = int(pad_len) + if cfg.survival_eps <= 0.0: + raise ValueError( + "fused top-k scheduling requires survival_eps > 0 so zeroed pad " + "rows cannot compete for the legacy policy budget" + ) + if not 0 <= num_real <= num_rows: + raise ValueError(f"num_real must be in [0, {num_rows}], got {num_real}") + max_token_len = int(cfg.resolved_max_verify_len) + 1 + if not 1 <= pad_len <= max_token_len: + raise ValueError(f"pad_len must be in [1, {max_token_len}], got {pad_len}") + + graph_num_tokens = int(graph_num_tokens) + token_floor = int(cfg.min_verify_len) + 1 + pad_tokens = (num_rows - num_real) * pad_len + minimum_tokens = num_real * token_floor + pad_tokens + maximum_tokens = num_real * max_token_len + pad_tokens + scheduled_capacity = min(max(budget, 0), num_real * int(cfg.schedulable_per_request)) + if not minimum_tokens + scheduled_capacity <= graph_num_tokens <= maximum_tokens: + raise ValueError( + "fused schedule cannot realize the selected CUDA graph without " + "changing policy semantics: " + f"expected [{minimum_tokens + scheduled_capacity}, {maximum_tokens}] " + f"tokens, graph requires {graph_num_tokens}" + ) + + if survival.is_cuda: + return _schedule_verify_lens_topk_fused_fill_triton( + survival=survival, + budget=budget, + num_real=num_real, + pad_len=pad_len, + graph_num_tokens=graph_num_tokens, + cfg=cfg, + out=out, + ) + + rows = torch.arange(num_rows, device=survival.device) + real_survival = torch.where( + (rows < num_real).unsqueeze(1), + survival, + torch.full_like(survival, -1.0), + ) + scheduled = schedule_verify_lens_topk( + survival=real_survival, + budget=budget, + cfg=cfg, + ) + token_lens = torch.where( + rows < num_real, + scheduled + 1, + torch.full_like(scheduled, pad_len), + ) + spare = graph_num_tokens - int(token_lens.sum()) + for _ in range(max_token_len - 1): + if spare <= 0: + break + headroom = (rows < num_real) & (token_lens < max_token_len) + grant = headroom & (torch.cumsum(headroom.to(torch.int64), dim=0) <= spare) + token_lens += grant.to(torch.int32) + spare -= int(grant.sum()) + if spare != 0: + raise ValueError(f"fused schedule left {spare} graph tokens after round-robin fill") + if out is None: + return token_lens + out[:num_rows].copy_(token_lens) + return out[:num_rows] diff --git a/tensorrt_llm/_torch/speculative/dspark_select.py b/tensorrt_llm/_torch/speculative/dspark_select.py new file mode 100644 index 000000000000..dff79b5f8454 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_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. +"""DSpark verify-window selection: the in-graph prologue. + +The host planner decides windows from a confidence snapshot that is one +iteration old (the batch reshuffles across the lag, and the read itself must +not sync). This module ranks the block that is about to be verified by its +OWN confidence instead: everything from the slot-indexed gather to the +per-token row maps is 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_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 8237c09ee8b7..ff345a77e7a4 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -1584,6 +1584,52 @@ def invalidate_sampling_params_cache(self) -> None: self._sampling_params_signature[0] = None self._sampling_params_signature[1] = None + def remap_expanded_sampling_params(self, req_idx: torch.Tensor, + num_tokens: int) -> None: + """Repack per-logits-row sampling state after device window selection. + + Device window selection can redistribute a fixed token budget among + requests immediately before replay. The expanded sampling buffers must + follow the resulting token owners. + """ + if num_tokens <= 0: + return + owners = req_idx[:num_tokens].to(torch.long) + pairs = ( + (self.request_temperatures, self.temperatures), + (self.request_top_ks, self.top_ks), + (self.request_top_ps, self.top_ps), + (self.request_min_ps, self.min_ps), + (self.request_seeds, self.seeds), + (self.request_offsets, self.offsets), + ) + for request_values, expanded_values in pairs: + if request_values is None or expanded_values is None: + continue + if expanded_values.numel() < num_tokens: + raise RuntimeError( + "device-selected DSpark layout exceeds the captured " + "sampling-parameter buffer") + torch.index_select(request_values, + 0, + owners, + out=expanded_values[:num_tokens]) + + if self.request_offsets is not None and self.offsets is not None: + if self.qo_indptr is None: + raise RuntimeError( + "device-selected DSpark sampling requires ragged row starts" + ) + row_slots = torch.arange(num_tokens, + device=owners.device, + dtype=torch.int64) + row_slots -= self.qo_indptr.index_select(0, owners) + self.offsets[:num_tokens].add_(row_slots * _RNG_SLOT_SPAN) + + # The expanded cache now follows the device-selected layout. Force the + # next host-prepared iteration to restore its own row expansion. + self._sampling_params_signature[1] = None + class AuxiliarySpeculativeStateHandler(Protocol): """Model-side state that verification writes once per draft token. diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index cf7ec9d4904b..5c819ec69042 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -158,6 +158,13 @@ def _get_comm_ranks_device_id(self): return comm_ranks, device_ids + def _get_pytorch_executor_factory(self): + """Keep normal construction unchanged; offline subclasses may supply a factory.""" + from tensorrt_llm._torch.pyexecutor.py_executor_creator import \ + create_py_executor + + return create_py_executor + def setup_engine(self): """ Setup the engine for the worker. @@ -173,9 +180,7 @@ def _create_py_executor(): ), "llm_args should be with backend in _create_py_executor" _ = self._get_comm_ranks_device_id() if self._backend == "pytorch": - from tensorrt_llm._torch.pyexecutor.py_executor_creator import \ - create_py_executor - create_executor = create_py_executor + create_executor = self._get_pytorch_executor_factory() args["llm_args"] = self.llm_args args["checkpoint_dir"] = self._hf_model_dir args["tokenizer"] = self._tokenizer diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index e8bbc442b234..0ab7de441515 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3583,6 +3583,32 @@ def spec_dec_mode(self): return TorchSpeculativeDecodingMode.DFLASH +class DSparkMixedEvidenceConfig(StrictBaseModel): + """Optional measured BREAKABLE mixed evidence, independent of decode costs. + + All SHA256 pins come from the authenticated deployment package. The live + fingerprint must be independently observed for the current runtime; copying + the cost table's identity does not qualify it. Uncovered or invalid mixed + evidence keeps full-width native verification without disabling admitted + decode confidence scheduling. + """ + + costs_path: str = Field( + description="Measured mixed geometry/action cost JSON.") + costs_sha256: str = Field(pattern=r"^[0-9a-f]{64}$", + description="Cost-file SHA256.") + admission_path: str = Field( + description="Qualified mixed admission receipt JSON.") + admission_sha256: str = Field(pattern=r"^[0-9a-f]{64}$", + description="Receipt SHA256.") + protocol_sha256: str = Field(pattern=r"^[0-9a-f]{64}$", + description="Measured protocol SHA256.") + live_fingerprint_path: str = Field( + description="Independently generated mixed runtime fingerprint JSON.") + live_fingerprint_sha256: str = Field(pattern=r"^[0-9a-f]{64}$", + description="Live fingerprint SHA256.") + + class DSparkDecodingConfig(DecodingBaseConfig): """Configuration for DSpark speculative decoding. @@ -3592,9 +3618,8 @@ class DSparkDecodingConfig(DecodingBaseConfig): backbone forward, but it additionally refines the per-position draft logits with a lightweight sequential head (a low-rank Markov head, optionally an RNN head). A confidence head can score each draft position for verification - scheduling without truncating the proposed block. This configuration selects - native full-block drafting and does not expose confidence-scheduling or - calibration controls. + scheduling without truncating the proposed block. Confidence scheduling is + opt-in; the default configuration retains native full-block verification. Key features: - Target-dependent: captures hidden states from ``target_layer_ids``. @@ -3636,6 +3661,89 @@ class DSparkDecodingConfig(DecodingBaseConfig): "read from the draft model config (dspark_markov_head_type), " "defaulting to \"vanilla\".") + enable_confidence_scheduling: bool = Field( + default=False, + description= + "Score each drafted position with the draft's confidence head and let " + "the verification scheduler size the batch's verify budget from those " + "scores. The block is always drafted in full; only the number of tokens " + "sent to the target for verification changes, so acceptance (and hence " + "output distribution) is unaffected. Requires a checkpoint with trained " + "confidence-head weights. Gains only appear at high concurrency -- at " + "small batch sizes the step cost is nearly flat in the verified token " + "count, so trimming saves nothing and costs acceptance length.", + status="prototype") + + validate_confidence_artifact_freshness: bool = Field( + default=False, + description= + "Require exact live cost-artifact identity and a positive workload " + "admission receipt before enabling confidence scheduling. By default, " + "explicitly enabled scheduling reuses structurally valid cost and " + "calibration artifacts, with a warning that performance predictions " + "may be stale. This never changes target verification or acceptance.", + status="prototype") + + confidence_sts_path: Optional[str] = Field( + default=None, + description= + "Path to a JSON file with per-position sequential-temperature-scaling " + "(STS) values used to calibrate the confidence head. The scheduler " + "consumes a cumulative product of per-position probabilities, so a " + "per-position calibration error compounds geometrically along the " + "block. If None, the raw sigmoid is used (identity calibration).", + status="prototype") + + confidence_sps_table_path: Optional[str] = Field( + default=None, + description= + "Path to an authenticated schema-v2 JSON table of measured decode step " + "costs keyed by exact rank-local CUDA-graph shape (G, V).", + status="prototype") + + confidence_sps_live_fingerprint_path: Optional[str] = Field( + default=None, + description= + "Path to an independently generated JSON fingerprint of the active " + "runtime. Required when validate_confidence_artifact_freshness=True. " + "Otherwise it is optional observed provenance, kept separate from " + "the cost artifact's original metadata.", + status="prototype") + + confidence_admission_receipt_path: Optional[str] = Field( + default=None, + description= + "Path to an authenticated workload-level confidence admission receipt. " + "Required only when validate_confidence_artifact_freshness=True; it " + "must prove positive aggregate compact value after the fixed " + "confidence-path overhead and safety margin are charged.", + status="prototype") + + confidence_admission_receipt_sha256: Optional[str] = Field( + default=None, + description="Externally authenticated lowercase SHA256 of " + "confidence_admission_receipt_path. This pin must come from the sealed " + "deployment package rather than from the receipt itself.", + status="prototype") + + enable_fused_confidence_scheduler: bool = Field( + default=False, + description= + "Use the optional policy-neutral Triton implementation of confidence " + "top-k allocation plus exact CUDA-graph bucket fill. This changes only " + "the scheduler launch path: budgets, survival thresholds, selected " + "windows, and fallback behavior remain identical. It is independently " + "default-off and requires confidence scheduling.", + status="prototype") + + confidence_mixed_evidence: Optional[DSparkMixedEvidenceConfig] = Field( + default=None, + description="Optional separately admitted physical-K3, eight-rank " + "BREAKABLE mixed scheduling evidence. Does not reuse decode T(G,V) costs. " + "None, unsupported routes, stale ownership, or uncovered mixed geometry " + "retain native mixed verification while admitted decode remains enabled.", + status="prototype") + decoding_type: Literal["DSpark"] = Field(default="DSpark") attention_backend: Literal["AUTO", "VANILLA", "TRTLLM", "CUTEDSL"] = Field( @@ -3651,6 +3759,51 @@ class DSparkDecodingConfig(DecodingBaseConfig): "and needs a cute-dsl MLA decode taking per-token kv_bounds that is not " "upstream yet. Which kernel each name selects: MLADSparkForCausalLM.") + @model_validator(mode="after") + def validate_confidence_scheduling(self): + if not self.enable_confidence_scheduling: + if (self.validate_confidence_artifact_freshness + or self.enable_fused_confidence_scheduler + or self.confidence_sts_path + or self.confidence_sps_table_path + or self.confidence_sps_live_fingerprint_path + or self.confidence_admission_receipt_path + or self.confidence_admission_receipt_sha256 + or self.confidence_mixed_evidence is not None): + raise ValueError( + "validate_confidence_artifact_freshness / " + "enable_fused_confidence_scheduler / confidence_sts_path / " + "confidence_sps_table_path / " + "confidence_sps_live_fingerprint_path / " + "confidence_admission_receipt_path / " + "confidence_admission_receipt_sha256 / confidence_mixed_evidence require " + "enable_confidence_scheduling=True") + return self + + if self.max_draft_len == 1: + # Physical K1 has no smaller verifier window, so confidence + # scheduling can only add control-path work. Resolve it to the + # ordinary feature-off configuration before any runtime component + # observes the config. This preserves the original static K path + # rather than running a confidence-enabled native fallback. + logger.warning( + "DSpark confidence scheduling has no compact tier below the " + "native K1 layout; using the feature-off static K1 path and " + "ignoring confidence-scheduling inputs.") + self.enable_confidence_scheduling = False + self.validate_confidence_artifact_freshness = False + self.enable_fused_confidence_scheduler = False + self.confidence_sts_path = None + self.confidence_sps_table_path = None + self.confidence_sps_live_fingerprint_path = None + self.confidence_admission_receipt_path = None + self.confidence_admission_receipt_sha256 = None + self.confidence_mixed_evidence = None + return self + # Resolve optional freshness admission before runtime components + # snapshot the flag. Exact graph/K coverage is checked by the engine. + return self + @model_validator(mode="after") def set_max_total_draft_tokens(self): self.max_total_draft_tokens = self.max_draft_len @@ -6900,6 +7053,170 @@ def validate_encode_only_torch_compile_config(self) -> 'TorchLlmArgs': "graphs or disable enable_piecewise_cuda_graph.") return self + def _validate_dspark_confidence_environment(self) -> None: + """Reject configurations confidence scheduling cannot honour. + + Every check here guards a path that would otherwise degrade *silently*: + the run comes up, produces plausible text, and either never takes the + ragged path at all or takes it with a kernel that mis-attributes rows. + A hard error at construction is the only signal that survives. + """ + if not self.speculative_config.draft_is_embedded_in_target: + raise ValueError( + "DSpark confidence scheduling supports only the embedded " + "DeepSeek-V4 draft. Standalone DSpark retains native " + "verification; disable confidence scheduling for that path.") + cuda_graph_config = self.cuda_graph_config + if self.pipeline_parallel_size > 1: + raise ValueError( + "speculative_config.enable_confidence_scheduling=True does not yet " + "support pipeline parallelism. The PP executor's pipelined " + "microbatches do not provide the single-iteration host-input " + "lifetime boundary required by ragged verification staging. " + "Set pipeline_parallel_size=1 or disable confidence scheduling." + ) + if self.guided_decoding_backend is not None: + raise ValueError( + "speculative_config.enable_confidence_scheduling=True does not yet " + "support guided decoding. Guided-decoder bitmasks are laid " + "out with one uniform speculative stride; applying them to " + "packed per-request verify windows would shift every request " + "after the first short window. Disable guided decoding or " + "disable confidence scheduling.") + if self.enable_lora or self.lora_config is not None: + raise ValueError( + "speculative_config.enable_confidence_scheduling=True does not yet " + "support LoRA. LoRA metadata expands generation requests " + "with one uniform speculative stride, which cannot represent " + "packed per-request verify windows. Disable LoRA or disable " + "confidence scheduling.") + # `fit_ragged_verify_lens` plans against the padded row count and + # reserves tokens for the CUDA-graph padding rows, so it presumes those + # rows exist. With padding off, `_get_padded_batch` adds none: the batch + # then carries a token total that keys into a graph nobody captured and + # drops to eager on every ragged step -- which costs far more than the + # trimmed tokens save -- while real requests were still shrunk to leave + # room for padding that never arrived. + if cuda_graph_config is None: + raise ValueError( + "speculative_config.enable_confidence_scheduling=True requires " + "cuda_graph_config to be set with enable_padding=True. Ragged " + "verification fits each step's token total onto a captured " + "CUDA-graph bucket; with no graphs there is no bucket to fit " + "and every step would silently run eager.") + if not cuda_graph_config.enable_padding: + raise ValueError( + "speculative_config.enable_confidence_scheduling=True requires " + "cuda_graph_config.enable_padding=True (it defaults to False). " + "The ragged token budget is planned against the padded batch " + "size and reserves tokens for the padding rows; without padding " + "those rows never appear, so the step's token total misses its " + "captured bucket and falls out of graph replay.") + + sparse_attention_config = self.sparse_attention_config + if sparse_attention_config is not None: + # Both recover "which request is this row" from a fixed next_n -- + # GVR indexes its previous-step top-k hint by request and strides it + # by next_n, and the CuTe-DSL kernel takes next_n as a compile-time + # constant and JIT cache key. Neither can express per-request + # windows. The C++ launcher already refuses the combination + # (indexerTopK.cu, "ragged rowKvLens is incompatible with the GVR + # preIdx hint"), but that fires deep in a decode step; say so here. + incompatible = [ + name for name in ("enable_heuristic_topk", "use_cute_dsl_topk") + if getattr(sparse_attention_config, name, False) + ] + if incompatible: + raise ValueError( + "speculative_config.enable_confidence_scheduling=True is " + "incompatible with sparse_attention_config." + f"{' / '.join(incompatible)}=True: those top-k paths " + "reconstruct each row's request and window position from a " + "single batch-wide next_n, which a ragged batch does not " + "have. Disable them, or disable confidence scheduling.") + # Not an error: the DSL paged-MQA-logits path is structurally + # uniform too, but the expanded (one row per query token) path it + # falls back to is numerically equivalent, and it is the default on + # DeepSeek-V4 anyway. Still worth saying out loud, since the knob + # will read as honoured otherwise. + if getattr(sparse_attention_config, "use_cute_dsl_paged_mqa_logits", + False): + logger.warning( + "sparse_attention_config.use_cute_dsl_paged_mqa_logits=True " + "is ignored on ragged-verification steps: the DSL kernel " + "broadcasts one KV length per request across a fixed next_n. " + "Those steps use the expanded DeepGEMM path instead, which " + "is numerically equivalent.") + + @staticmethod + def _resolve_dspark_confidence_workload_admission( + spec_cfg: DSparkDecodingConfig) -> None: + """Resolve confidence activation before engines or workers observe it.""" + from tensorrt_llm._torch.speculative.dspark_planner import \ + load_confidence_workload_admission + + receipt_path = spec_cfg.confidence_admission_receipt_path + receipt_sha256 = spec_cfg.confidence_admission_receipt_sha256 + sps_table_path = spec_cfg.confidence_sps_table_path + live_fingerprint_path = spec_cfg.confidence_sps_live_fingerprint_path + if not spec_cfg.validate_confidence_artifact_freshness: + if not sps_table_path: + raise ValueError( + "DSpark confidence scheduling requires confidence_sps_table_path; " + "a missing measured cost table cannot be timed as enabled.") + logger.warning( + "DSpark confidence scheduling is enabled with reusable artifacts. " + "Cost/calibration predictions and performance economics are not " + "freshness-validated and may be stale; target acceptance is unchanged. " + "Set validate_confidence_artifact_freshness=True to require exact " + "live provenance and positive workload admission.") + return + missing = [ + name for name, value in ( + ("confidence_admission_receipt_path", receipt_path), + ("confidence_admission_receipt_sha256", receipt_sha256), + ("confidence_sps_table_path", sps_table_path), + ("confidence_sps_live_fingerprint_path", live_fingerprint_path), + ) if not value + ] + fallback_reason = ("missing " + ", ".join(missing)) if missing else None + if fallback_reason is None: + assert receipt_path is not None + assert receipt_sha256 is not None + assert sps_table_path is not None + assert live_fingerprint_path is not None + try: + admission = load_confidence_workload_admission( + receipt_path, + expected_receipt_sha256=receipt_sha256, + sps_cost_table_path=sps_table_path, + live_engine_fingerprint_path=live_fingerprint_path, + physical_k=spec_cfg.max_draft_len, + ) + if not admission.admitted: + fallback_reason = ( + "the authenticated workload replay does not have positive " + "net compact value after fixed overhead and safety margin" + ) + except (OSError, TypeError, ValueError) as error: + fallback_reason = str(error) + + if fallback_reason is None: + return + + logger.warning( + f"DSpark confidence scheduling is not deployment-admitted ({fallback_reason}); " + f"using the feature-off static K{spec_cfg.max_draft_len} path.") + spec_cfg.enable_confidence_scheduling = False + spec_cfg.validate_confidence_artifact_freshness = False + spec_cfg.enable_fused_confidence_scheduler = False + spec_cfg.confidence_sts_path = None + spec_cfg.confidence_sps_table_path = None + spec_cfg.confidence_sps_live_fingerprint_path = None + spec_cfg.confidence_admission_receipt_path = None + spec_cfg.confidence_admission_receipt_sha256 = None + spec_cfg.confidence_mixed_evidence = None + @model_validator(mode="after") def normalize_prefill_cuda_graph_config(self) -> 'TorchLlmArgs': """Normalize legacy piecewise CUDA graph options into prefill fields.""" @@ -7379,6 +7696,10 @@ def _dspark_get( "DSpark block_size must equal max_draft_len; got " f"block_size={spec_cfg.block_size} and " f"max_draft_len={spec_cfg.max_draft_len}") + if spec_cfg.enable_confidence_scheduling: + self._resolve_dspark_confidence_workload_admission(spec_cfg) + if spec_cfg.enable_confidence_scheduling: + self._validate_dspark_confidence_environment() if isinstance(self.speculative_config, SADecodingConfig): pool_size = self.speculative_config.global_pool_size diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 1211e43a6390..c5b5bc4f3294 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -930,7 +930,7 @@ }, "DSparkDecodingConfig": { "additionalProperties": false, - "description": "Configuration for DSpark speculative decoding.\n\nDSpark (DeepSeek) is a target-dependent, \"semi-parallel\" speculative\ndecoding method. Like DFlash it captures hidden states from several target\nlayers as cross-attention context and drafts a whole block in a single\nbackbone forward, but it additionally refines the per-position draft logits\nwith a lightweight sequential head (a low-rank Markov head, optionally an RNN\nhead). A confidence head can score each draft position for verification\nscheduling without truncating the proposed block. This configuration selects\nnative full-block drafting and does not expose confidence-scheduling or\ncalibration controls.\n\nKey features:\n- Target-dependent: captures hidden states from ``target_layer_ids``.\n- Semi-parallel: one block backbone forward + cheap sequential head refine.\n- Full-block proposals: confidence scoring does not change the draft tokens\n or the standard target-verification acceptance rule.\n\nReference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec).", + "description": "Configuration for DSpark speculative decoding.\n\nDSpark (DeepSeek) is a target-dependent, \"semi-parallel\" speculative\ndecoding method. Like DFlash it captures hidden states from several target\nlayers as cross-attention context and drafts a whole block in a single\nbackbone forward, but it additionally refines the per-position draft logits\nwith a lightweight sequential head (a low-rank Markov head, optionally an RNN\nhead). A confidence head can score each draft position for verification\nscheduling without truncating the proposed block. Confidence scheduling is\nopt-in; the default configuration retains native full-block verification.\n\nKey features:\n- Target-dependent: captures hidden states from ``target_layer_ids``.\n- Semi-parallel: one block backbone forward + cheap sequential head refine.\n- Full-block proposals: confidence scoring does not change the draft tokens\n or the standard target-verification acceptance rule.\n\nReference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec).", "properties": { "acceptance_rate_threshold": { "anyOf": [ @@ -999,6 +999,89 @@ "description": "Number of draft positions produced per block. If None, read from the draft model config (dspark_block_size). Should equal max_draft_len.", "title": "Block Size" }, + "confidence_admission_receipt_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to an authenticated workload-level confidence admission receipt. Required only when validate_confidence_artifact_freshness=True; it must prove positive aggregate compact value after the fixed confidence-path overhead and safety margin are charged.", + "status": "prototype", + "title": "Confidence Admission Receipt Path" + }, + "confidence_admission_receipt_sha256": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Externally authenticated lowercase SHA256 of confidence_admission_receipt_path. This pin must come from the sealed deployment package rather than from the receipt itself.", + "status": "prototype", + "title": "Confidence Admission Receipt Sha256" + }, + "confidence_mixed_evidence": { + "anyOf": [ + { + "$ref": "#/$defs/DSparkMixedEvidenceConfig" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Optional separately admitted physical-K3, eight-rank BREAKABLE mixed scheduling evidence. Does not reuse decode T(G,V) costs. None, unsupported routes, stale ownership, or uncovered mixed geometry retain native mixed verification while admitted decode remains enabled.", + "status": "prototype" + }, + "confidence_sps_live_fingerprint_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to an independently generated JSON fingerprint of the active runtime. Required when validate_confidence_artifact_freshness=True. Otherwise it is optional observed provenance, kept separate from the cost artifact's original metadata.", + "status": "prototype", + "title": "Confidence Sps Live Fingerprint Path" + }, + "confidence_sps_table_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to an authenticated schema-v2 JSON table of measured decode step costs keyed by exact rank-local CUDA-graph shape (G, V).", + "status": "prototype", + "title": "Confidence Sps Table Path" + }, + "confidence_sts_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to a JSON file with per-position sequential-temperature-scaling (STS) values used to calibrate the confidence head. The scheduler consumes a cumulative product of per-position probabilities, so a per-position calibration error compounds geometrically along the block. If None, the raw sigmoid is used (identity calibration).", + "status": "prototype", + "title": "Confidence Sts Path" + }, "decoding_type": { "const": "DSpark", "default": "DSpark", @@ -1021,6 +1104,20 @@ "description": "Developer interface: dynamically adjust draft length based on active batch size in runtime.Maps batch size to draft lengths.For example: draft_len_schedule = {4:4, 8:2, 32:1} - Batch sizes 1-4: use draft_len=4 - Batch sizes 5-8: use draft_len=2 - Batch sizes 9-32: use draft_len=1 - Batch sizes 33+: use draft_len=0 (implicit, speculation disabled). Mutually exclusive with max_concurrency since draft_len_schedule implicitly support max concurrency control.", "title": "Draft Len Schedule" }, + "enable_confidence_scheduling": { + "default": false, + "description": "Score each drafted position with the draft's confidence head and let the verification scheduler size the batch's verify budget from those scores. The block is always drafted in full; only the number of tokens sent to the target for verification changes, so acceptance (and hence output distribution) is unaffected. Requires a checkpoint with trained confidence-head weights. Gains only appear at high concurrency -- at small batch sizes the step cost is nearly flat in the verified token count, so trimming saves nothing and costs acceptance length.", + "status": "prototype", + "title": "Enable Confidence Scheduling", + "type": "boolean" + }, + "enable_fused_confidence_scheduler": { + "default": false, + "description": "Use the optional policy-neutral Triton implementation of confidence top-k allocation plus exact CUDA-graph bucket fill. This changes only the scheduler launch path: budgets, survival thresholds, selected windows, and fallback behavior remain identical. It is independently default-off and requires confidence scheduling.", + "status": "prototype", + "title": "Enable Fused Confidence Scheduler", + "type": "boolean" + }, "enable_penalty": { "default": false, "description": "If true, enables the occurrence penalties (repetition / presence / frequency) for one-model speculative decoding. Off by default because the penalties need a [num_seq_slots, vocab_size] occurrence-count workspace that is allocated up front (CUDA graphs capture fixed buffer addresses). While off, a request that asks for any of these penalties is rejected at admission rather than silently decoded without them.", @@ -1209,11 +1306,74 @@ "status": "prototype", "title": "Use Rejection Sampling", "type": "boolean" + }, + "validate_confidence_artifact_freshness": { + "default": false, + "description": "Require exact live cost-artifact identity and a positive workload admission receipt before enabling confidence scheduling. By default, explicitly enabled scheduling reuses structurally valid cost and calibration artifacts, with a warning that performance predictions may be stale. This never changes target verification or acceptance.", + "status": "prototype", + "title": "Validate Confidence Artifact Freshness", + "type": "boolean" } }, "title": "DSparkDecodingConfig", "type": "object" }, + "DSparkMixedEvidenceConfig": { + "additionalProperties": false, + "description": "Optional measured BREAKABLE mixed evidence, independent of decode costs.\n\nAll SHA256 pins come from the authenticated deployment package. The live\nfingerprint must be independently observed for the current runtime; copying\nthe cost table's identity does not qualify it. Uncovered or invalid mixed\nevidence keeps full-width native verification without disabling admitted\ndecode confidence scheduling.", + "properties": { + "admission_path": { + "description": "Qualified mixed admission receipt JSON.", + "title": "Admission Path", + "type": "string" + }, + "admission_sha256": { + "description": "Receipt SHA256.", + "pattern": "^[0-9a-f]{64}$", + "title": "Admission Sha256", + "type": "string" + }, + "costs_path": { + "description": "Measured mixed geometry/action cost JSON.", + "title": "Costs Path", + "type": "string" + }, + "costs_sha256": { + "description": "Cost-file SHA256.", + "pattern": "^[0-9a-f]{64}$", + "title": "Costs Sha256", + "type": "string" + }, + "live_fingerprint_path": { + "description": "Independently generated mixed runtime fingerprint JSON.", + "title": "Live Fingerprint Path", + "type": "string" + }, + "live_fingerprint_sha256": { + "description": "Live fingerprint SHA256.", + "pattern": "^[0-9a-f]{64}$", + "title": "Live Fingerprint Sha256", + "type": "string" + }, + "protocol_sha256": { + "description": "Measured protocol SHA256.", + "pattern": "^[0-9a-f]{64}$", + "title": "Protocol Sha256", + "type": "string" + } + }, + "required": [ + "costs_path", + "costs_sha256", + "admission_path", + "admission_sha256", + "protocol_sha256", + "live_fingerprint_path", + "live_fingerprint_sha256" + ], + "title": "DSparkMixedEvidenceConfig", + "type": "object" + }, "DecodeCudaGraphConfig": { "additionalProperties": false, "description": "CUDA graph configuration for decode requests.", diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index 15599ad03a1f..d56c50104248 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -930,7 +930,7 @@ }, "DSparkDecodingConfig": { "additionalProperties": false, - "description": "Configuration for DSpark speculative decoding.\n\nDSpark (DeepSeek) is a target-dependent, \"semi-parallel\" speculative\ndecoding method. Like DFlash it captures hidden states from several target\nlayers as cross-attention context and drafts a whole block in a single\nbackbone forward, but it additionally refines the per-position draft logits\nwith a lightweight sequential head (a low-rank Markov head, optionally an RNN\nhead). A confidence head can score each draft position for verification\nscheduling without truncating the proposed block. This configuration selects\nnative full-block drafting and does not expose confidence-scheduling or\ncalibration controls.\n\nKey features:\n- Target-dependent: captures hidden states from ``target_layer_ids``.\n- Semi-parallel: one block backbone forward + cheap sequential head refine.\n- Full-block proposals: confidence scoring does not change the draft tokens\n or the standard target-verification acceptance rule.\n\nReference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec).", + "description": "Configuration for DSpark speculative decoding.\n\nDSpark (DeepSeek) is a target-dependent, \"semi-parallel\" speculative\ndecoding method. Like DFlash it captures hidden states from several target\nlayers as cross-attention context and drafts a whole block in a single\nbackbone forward, but it additionally refines the per-position draft logits\nwith a lightweight sequential head (a low-rank Markov head, optionally an RNN\nhead). A confidence head can score each draft position for verification\nscheduling without truncating the proposed block. Confidence scheduling is\nopt-in; the default configuration retains native full-block verification.\n\nKey features:\n- Target-dependent: captures hidden states from ``target_layer_ids``.\n- Semi-parallel: one block backbone forward + cheap sequential head refine.\n- Full-block proposals: confidence scoring does not change the draft tokens\n or the standard target-verification acceptance rule.\n\nReference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec).", "properties": { "acceptance_rate_threshold": { "anyOf": [ @@ -999,6 +999,89 @@ "description": "Number of draft positions produced per block. If None, read from the draft model config (dspark_block_size). Should equal max_draft_len.", "title": "Block Size" }, + "confidence_admission_receipt_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to an authenticated workload-level confidence admission receipt. Required only when validate_confidence_artifact_freshness=True; it must prove positive aggregate compact value after the fixed confidence-path overhead and safety margin are charged.", + "status": "prototype", + "title": "Confidence Admission Receipt Path" + }, + "confidence_admission_receipt_sha256": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Externally authenticated lowercase SHA256 of confidence_admission_receipt_path. This pin must come from the sealed deployment package rather than from the receipt itself.", + "status": "prototype", + "title": "Confidence Admission Receipt Sha256" + }, + "confidence_mixed_evidence": { + "anyOf": [ + { + "$ref": "#/$defs/DSparkMixedEvidenceConfig" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Optional separately admitted physical-K3, eight-rank BREAKABLE mixed scheduling evidence. Does not reuse decode T(G,V) costs. None, unsupported routes, stale ownership, or uncovered mixed geometry retain native mixed verification while admitted decode remains enabled.", + "status": "prototype" + }, + "confidence_sps_live_fingerprint_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to an independently generated JSON fingerprint of the active runtime. Required when validate_confidence_artifact_freshness=True. Otherwise it is optional observed provenance, kept separate from the cost artifact's original metadata.", + "status": "prototype", + "title": "Confidence Sps Live Fingerprint Path" + }, + "confidence_sps_table_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to an authenticated schema-v2 JSON table of measured decode step costs keyed by exact rank-local CUDA-graph shape (G, V).", + "status": "prototype", + "title": "Confidence Sps Table Path" + }, + "confidence_sts_path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Path to a JSON file with per-position sequential-temperature-scaling (STS) values used to calibrate the confidence head. The scheduler consumes a cumulative product of per-position probabilities, so a per-position calibration error compounds geometrically along the block. If None, the raw sigmoid is used (identity calibration).", + "status": "prototype", + "title": "Confidence Sts Path" + }, "decoding_type": { "const": "DSpark", "default": "DSpark", @@ -1021,6 +1104,20 @@ "description": "Developer interface: dynamically adjust draft length based on active batch size in runtime.Maps batch size to draft lengths.For example: draft_len_schedule = {4:4, 8:2, 32:1} - Batch sizes 1-4: use draft_len=4 - Batch sizes 5-8: use draft_len=2 - Batch sizes 9-32: use draft_len=1 - Batch sizes 33+: use draft_len=0 (implicit, speculation disabled). Mutually exclusive with max_concurrency since draft_len_schedule implicitly support max concurrency control.", "title": "Draft Len Schedule" }, + "enable_confidence_scheduling": { + "default": false, + "description": "Score each drafted position with the draft's confidence head and let the verification scheduler size the batch's verify budget from those scores. The block is always drafted in full; only the number of tokens sent to the target for verification changes, so acceptance (and hence output distribution) is unaffected. Requires a checkpoint with trained confidence-head weights. Gains only appear at high concurrency -- at small batch sizes the step cost is nearly flat in the verified token count, so trimming saves nothing and costs acceptance length.", + "status": "prototype", + "title": "Enable Confidence Scheduling", + "type": "boolean" + }, + "enable_fused_confidence_scheduler": { + "default": false, + "description": "Use the optional policy-neutral Triton implementation of confidence top-k allocation plus exact CUDA-graph bucket fill. This changes only the scheduler launch path: budgets, survival thresholds, selected windows, and fallback behavior remain identical. It is independently default-off and requires confidence scheduling.", + "status": "prototype", + "title": "Enable Fused Confidence Scheduler", + "type": "boolean" + }, "enable_penalty": { "default": false, "description": "If true, enables the occurrence penalties (repetition / presence / frequency) for one-model speculative decoding. Off by default because the penalties need a [num_seq_slots, vocab_size] occurrence-count workspace that is allocated up front (CUDA graphs capture fixed buffer addresses). While off, a request that asks for any of these penalties is rejected at admission rather than silently decoded without them.", @@ -1209,11 +1306,74 @@ "status": "prototype", "title": "Use Rejection Sampling", "type": "boolean" + }, + "validate_confidence_artifact_freshness": { + "default": false, + "description": "Require exact live cost-artifact identity and a positive workload admission receipt before enabling confidence scheduling. By default, explicitly enabled scheduling reuses structurally valid cost and calibration artifacts, with a warning that performance predictions may be stale. This never changes target verification or acceptance.", + "status": "prototype", + "title": "Validate Confidence Artifact Freshness", + "type": "boolean" } }, "title": "DSparkDecodingConfig", "type": "object" }, + "DSparkMixedEvidenceConfig": { + "additionalProperties": false, + "description": "Optional measured BREAKABLE mixed evidence, independent of decode costs.\n\nAll SHA256 pins come from the authenticated deployment package. The live\nfingerprint must be independently observed for the current runtime; copying\nthe cost table's identity does not qualify it. Uncovered or invalid mixed\nevidence keeps full-width native verification without disabling admitted\ndecode confidence scheduling.", + "properties": { + "admission_path": { + "description": "Qualified mixed admission receipt JSON.", + "title": "Admission Path", + "type": "string" + }, + "admission_sha256": { + "description": "Receipt SHA256.", + "pattern": "^[0-9a-f]{64}$", + "title": "Admission Sha256", + "type": "string" + }, + "costs_path": { + "description": "Measured mixed geometry/action cost JSON.", + "title": "Costs Path", + "type": "string" + }, + "costs_sha256": { + "description": "Cost-file SHA256.", + "pattern": "^[0-9a-f]{64}$", + "title": "Costs Sha256", + "type": "string" + }, + "live_fingerprint_path": { + "description": "Independently generated mixed runtime fingerprint JSON.", + "title": "Live Fingerprint Path", + "type": "string" + }, + "live_fingerprint_sha256": { + "description": "Live fingerprint SHA256.", + "pattern": "^[0-9a-f]{64}$", + "title": "Live Fingerprint Sha256", + "type": "string" + }, + "protocol_sha256": { + "description": "Measured protocol SHA256.", + "pattern": "^[0-9a-f]{64}$", + "title": "Protocol Sha256", + "type": "string" + } + }, + "required": [ + "costs_path", + "costs_sha256", + "admission_path", + "admission_sha256", + "protocol_sha256", + "live_fingerprint_path", + "live_fingerprint_sha256" + ], + "title": "DSparkMixedEvidenceConfig", + "type": "object" + }, "DecodeCudaGraphConfig": { "additionalProperties": false, "description": "CUDA graph configuration for decode requests.", diff --git a/tensorrt_llm/tools/dspark_calibration.py b/tensorrt_llm/tools/dspark_calibration.py new file mode 100644 index 000000000000..e705922cb234 --- /dev/null +++ b/tensorrt_llm/tools/dspark_calibration.py @@ -0,0 +1,239 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Offline chain-confidence STS fitting from correctly paired observation shards. + +This CPU tool imports PyTorch, not TensorRT-LLM or CUDA. Collection is a separate +runtime adapter: each row must join raw logits to full-block prefix-accept labels +from the same stamped draft sequence. Calibration is not task qualification or +an economic admission receipt. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +from pathlib import Path + +import torch + +__all__ = [ + "default_temperature_grid", + "expected_calibration_error", + "fit_sts_temperatures", + "load_shards", +] + +_EPS_PROB = 1e-8 + + +def default_temperature_grid() -> torch.Tensor: + """Return the established 41-point log-spaced STS grid over [0.1, 10].""" + return torch.logspace(math.log10(0.1), math.log10(10.0), steps=41) + + +def expected_calibration_error( + *, probs: torch.Tensor, targets: torch.Tensor, num_bins: int +) -> float: + """Compute population-weighted cumulative-prefix calibration error.""" + probs = probs.reshape(-1).to(torch.float64).clamp(_EPS_PROB, 1.0 - _EPS_PROB) + targets = targets.reshape(-1).to(torch.float64) + total = probs.numel() + if total == 0: + return float("nan") + bin_index = (probs * num_bins).long().clamp_(0, num_bins - 1) + count = torch.zeros(num_bins, dtype=torch.float64) + pred_sum = torch.zeros(num_bins, dtype=torch.float64) + target_sum = torch.zeros(num_bins, dtype=torch.float64) + count.scatter_add_(0, bin_index, torch.ones_like(probs)) + pred_sum.scatter_add_(0, bin_index, probs) + target_sum.scatter_add_(0, bin_index, targets) + denom = count.clamp_min(1.0) + bin_error = (pred_sum / denom - target_sum / denom).abs() + return float((bin_error * count).sum().item() / total) + + +def _validate_rows(logits: torch.Tensor, prefix_mask: torch.Tensor) -> None: + if ( + not isinstance(logits, torch.Tensor) + or not isinstance(prefix_mask, torch.Tensor) + or logits.device.type != "cpu" + or prefix_mask.device.type != "cpu" + or not logits.is_floating_point() + or logits.ndim != 2 + or logits.shape != prefix_mask.shape + or logits.shape[0] == 0 + or logits.shape[1] == 0 + ): + raise ValueError("CPU logits and prefix_mask must have identical nonempty [N, K] shapes") + if not torch.isfinite(logits).all() or not torch.isfinite(prefix_mask).all(): + raise ValueError("Calibration observations must be finite") + if not ((prefix_mask == 0) | (prefix_mask == 1)).all(): + raise ValueError("Prefix labels must be binary full-block acceptance observations") + if (prefix_mask[:, 1:] > prefix_mask[:, :-1]).any(): + raise ValueError("A later accepted prefix cannot follow an earlier rejected prefix") + + +def fit_sts_temperatures( + *, logits: torch.Tensor, prefix_mask: torch.Tensor, grid: torch.Tensor, num_bins: int = 15 +) -> dict[str, list[float]]: + """Fit each temperature left-to-right against cumulative survival ECE. + + Args: + logits: Raw confidence logits, CPU tensor [N, K]. + prefix_mask: Binary cumulative-prefix labels, CPU tensor [N, K]. + grid: Positive finite temperature candidates, CPU tensor [M]. + num_bins: Positive number of calibration-error bins. + """ + _validate_rows(logits, prefix_mask) + if type(num_bins) is not int or num_bins < 1: + raise ValueError("num_bins must be a positive integer") + if ( + not isinstance(grid, torch.Tensor) + or grid.device.type != "cpu" + or grid.ndim != 1 + or grid.numel() == 0 + or not torch.isfinite(grid).all() + or not (grid > 0).all() + ): + raise ValueError("The temperature grid must be a nonempty positive finite CPU vector") + logits = logits.to(torch.float64) + prefix_mask = prefix_mask.to(torch.float64) + num_samples, block_size = logits.shape + grid_values = grid.to(torch.float64).tolist() + temperatures, ece_before, ece_after = [], [], [] + survival_at_one = torch.ones(num_samples, dtype=torch.float64) + survival_fitted = torch.ones(num_samples, dtype=torch.float64) + for position in range(block_size): + position_logits = logits[:, position] + position_target = prefix_mask[:, position] + survival_at_one = survival_at_one * torch.sigmoid(position_logits) + ece_before.append( + expected_calibration_error( + probs=survival_at_one, targets=position_target, num_bins=num_bins + ) + ) + best_t = grid_values[0] + best_survival = survival_fitted * torch.sigmoid(position_logits / best_t) + best_ece = expected_calibration_error( + probs=best_survival, targets=position_target, num_bins=num_bins + ) + for temperature in grid_values[1:]: + candidate = survival_fitted * torch.sigmoid(position_logits / temperature) + ece = expected_calibration_error( + probs=candidate, targets=position_target, num_bins=num_bins + ) + if ece < best_ece: + best_ece, best_t, best_survival = ece, temperature, candidate + temperatures.append(float(best_t)) + ece_after.append(float(best_ece)) + survival_fitted = best_survival + return {"temperatures": temperatures, "ece_before": ece_before, "ece_after": ece_after} + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while block := stream.read(1 << 20): + digest.update(block) + return digest.hexdigest() + + +def load_shards( + paths: list[Path], physical_k: int +) -> tuple[torch.Tensor, torch.Tensor, dict[str, str]]: + """Load safe CPU tensor shards without conflating widths or label pairings. + + Pairing metadata is a producer contract, not independent proof of the + running checkpoint or label provenance. The original shard hashes remain + recorded; source or GPU freshness is not rewritten or certified. + """ + if type(physical_k) is not int or physical_k < 1 or not paths: + raise ValueError("An explicit positive physical K and input shards are required") + canonical = [path.resolve(strict=True) for path in paths] + if len(set(canonical)) != len(canonical): + raise ValueError("Duplicate calibration shards would double-count samples") + logits, masks, hashes = [], [], {} + for path, resolved in zip(paths, canonical): + if path.is_symlink() or not resolved.is_file(): + raise ValueError("Calibration inputs must be regular, nonsymlink files") + before = resolved.stat() + pin = _sha256_file(resolved) + blob = torch.load(resolved, map_location="cpu", weights_only=True) + if not isinstance(blob, dict) or not isinstance(blob.get("meta"), dict): + raise ValueError("Calibration shards require explicit pairing metadata") + if blob["meta"].get("pairing") != "draft_seq_ring": + raise ValueError("Recollect labels paired with the same stamped draft sequence") + raw_logits, prefix_mask = blob.get("logits"), blob.get("prefix_mask") + _validate_rows(raw_logits, prefix_mask) + if raw_logits.shape[1] != physical_k: + raise ValueError("The calibration shard width does not match the explicit physical K") + after = resolved.stat() + if (before.st_dev, before.st_ino, before.st_size, before.st_mtime_ns) != ( + after.st_dev, + after.st_ino, + after.st_size, + after.st_mtime_ns, + ) or _sha256_file(resolved) != pin: + raise ValueError("Calibration shard changed during reading") + logits.append(raw_logits) + masks.append(prefix_mask) + hashes[str(path)] = pin + return torch.cat(logits, 0), torch.cat(masks, 0), hashes + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, allow_abbrev=False) + parser.add_argument("--shards", nargs="+", type=Path, required=True) + parser.add_argument("--physical-k", type=int, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--num-bins", type=int, default=15) + parser.add_argument("--minimum-samples", type=int, default=1000) + args = parser.parse_args(argv) + if args.minimum_samples < 1: + raise ValueError("minimum-samples must be positive") + logits, masks, hashes = load_shards(args.shards, args.physical_k) + if logits.shape[0] < args.minimum_samples: + raise ValueError("Not enough paired observations to fit the requested calibration") + result = fit_sts_temperatures( + logits=logits, prefix_mask=masks, grid=default_temperature_grid(), num_bins=args.num_bins + ) + rates = masks.to(torch.float64).mean(dim=0).tolist() + payload = { + "sts_temperatures": result["temperatures"], + **result, + "physical_k": args.physical_k, + "num_samples": int(logits.shape[0]), + "num_bins": args.num_bins, + "raw_shard_sha256": hashes, + "pairing": "draft_seq_ring", + "prefix_acceptance_rate": rates, + "degenerate_positions": [index for index, rate in enumerate(rates) if rate in (0.0, 1.0)], + "objective": "left-to-right cumulative-prefix ECE;41-point temperature grid", + "producer_pairing_independently_verified": False, + "task_accuracy_or_economic_admission_qualified": False, + "artifact_freshness_validated": False, + } + if args.output.is_symlink(): + raise ValueError("Calibration output cannot be a symlink") + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0) + with os.fdopen(os.open(args.output, flags, 0o600), "w", encoding="utf-8") as stream: + json.dump(payload, stream, sort_keys=True, allow_nan=False) + stream.write("\n") + print( + json.dumps( + { + "output": str(args.output), + "num_samples": payload["num_samples"], + "physical_k": args.physical_k, + "degenerate_positions": payload["degenerate_positions"], + } + ) + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tensorrt_llm/tools/dspark_cost_export.py b/tensorrt_llm/tools/dspark_cost_export.py new file mode 100644 index 000000000000..b14e153ae409 --- /dev/null +++ b/tensorrt_llm/tools/dspark_cost_export.py @@ -0,0 +1,413 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Export direct DSpark whole-step costs or check existing artifacts. + +This standalone command uses the actual schema-v2 consumer. Only the explicit +complete-loop producer can export new costs; graph replay timings stay separate. +Run it by file path to avoid importing the native TensorRT-LLM package on CPU. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.util +import json +import os +import re +import statistics +import subprocess +import sys +from pathlib import Path +from types import ModuleType + +__all__ = [ + "check_cost_bytes", + "main", + "observed_engine_fingerprint", + "recorded_source_origin", + "source_origin_from_args", + "source_origin_input_metadata", + "runtime_cost_bytes", +] + +WHOLE_STEP_SCOPE = "original-executor-loop-hostStepTimeMS" + + +def _canonical_sha256(value: object) -> str: + return hashlib.sha256( + json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() + ).hexdigest() + + +def source_origin_input_metadata( + recorded_origin: Path | None, recorded_origin_sha256: str | None +) -> dict[str, str | bool]: + """Label recorded provenance separately from a fresh Git observation.""" + if (recorded_origin is None) != (recorded_origin_sha256 is None): + raise ValueError("Recorded source origin requires both its path and SHA256") + if recorded_origin is None: + return {"method": "git", "recorded_input": False} + if not isinstance(recorded_origin_sha256, str) or not re.fullmatch( + "[0-9a-f]{64}", recorded_origin_sha256 + ): + raise ValueError("Recorded source origin SHA256 must be explicit lowercase hex") + return { + "method": "recorded_provenance_input", + "recorded_input": True, + "input_sha256": recorded_origin_sha256, + "fresh_git_diff_repeated": False, + } + + +def source_origin_from_args(args: argparse.Namespace) -> dict[str, str]: + """Use one shared optional provenance input across the public drivers.""" + return recorded_source_origin( + args.source_root, + recorded_origin=getattr(args, "recorded_source_origin", None), + recorded_origin_sha256=getattr(args, "recorded_source_origin_sha256", None), + ) + + +def recorded_source_origin( + source_root: Path, + *, + recorded_origin: Path | None = None, + recorded_origin_sha256: str | None = None, +) -> dict[str, str]: + """Read Git provenance or an explicitly pinned record for this exact checkout. + + An optional record contains actual earlier HEAD/full-diff observation for + the same immutable source root. It is recorded metadata, not proof of source + freshness, GPU health, numerical accuracy or performance admission. Loaded + module checks and caller-owned source byte validation remain independent. + """ + source_origin_input_metadata(recorded_origin, recorded_origin_sha256) + if recorded_origin is not None: + if recorded_origin.is_symlink() or not recorded_origin.is_file(): + raise ValueError("Recorded source origin must be a regular non-symlink file") + if recorded_origin.stat().st_size > 8192: + raise ValueError("Recorded source origin must be a finite metadata record") + data = recorded_origin.read_bytes() + if hashlib.sha256(data).hexdigest() != recorded_origin_sha256: + raise ValueError("Recorded source origin bytes differ from the supplied SHA256") + utility = _source_module( + "_dspark_recorded_origin_json", Path(__file__).with_name("verify_costs.py") + ) + record = utility.parse_json(data) + if ( + not isinstance(record, dict) + or set(record) != {"schema", "source_root", "origin"} + or record["schema"] != "trtllm-recorded-source-origin-v1" + or record["source_root"] != str(source_root.resolve(strict=True)) + or not source_root.is_dir() + ): + raise ValueError("Recorded source origin belongs to a different checkout or schema") + origin = record["origin"] + if ( + not isinstance(origin, dict) + or set(origin) != {"source_head", "source_diff_sha256"} + or not isinstance(origin["source_head"], str) + or not re.fullmatch("[0-9a-f]{40}", origin["source_head"]) + or not isinstance(origin["source_diff_sha256"], str) + or not re.fullmatch("[0-9a-f]{64}", origin["source_diff_sha256"]) + ): + raise ValueError("Recorded source HEAD/full-diff digests are malformed") + return origin + head = ( + subprocess.run( + ["git", "--no-optional-locks", "-C", str(source_root), "rev-parse", "HEAD"], + check=True, + capture_output=True, + timeout=30, + ) + .stdout.decode() + .strip() + ) + diff = subprocess.run( + [ + "git", + "--no-optional-locks", + "-C", + str(source_root), + "diff", + "--no-ext-diff", + "--no-textconv", + "--binary", + "HEAD", + "--", + ], + check=True, + capture_output=True, + timeout=30, + ).stdout + return {"source_head": head, "source_diff_sha256": hashlib.sha256(diff).hexdigest()} + + +def observed_engine_fingerprint( + identity: dict[str, object], + *, + source_root: Path, + world_size: int, + recorded_origin: Path | None = None, + recorded_origin_sha256: str | None = None, +) -> dict[str, object]: + """Record actual origin metadata, not a destination freshness certificate. + + Loaded module bytes must be those of the supplied public checkout. Native + and runtime snapshots preserve observed library/configuration identity; no + old artifact fingerprint is relabeled as current and no GPU health is implied. + """ + topology = identity["topology"] + if world_size > 1 and not topology["attention_dp"]: + raise ValueError( + "Schema-v2 global graph counts use attention-DP fanout; " + "multi-rank attention-DP-off costs need a separate schema decision" + ) + sources = identity["loaded_source_sha256"] + for module, expected in sources.items(): + path = source_root / (module.replace(".", "/") + ".py") + if not path.is_file() or hashlib.sha256(path.read_bytes()).hexdigest() != expected: + raise ValueError("Loaded public source differs from the recorded origin checkout") + devices = identity["devices"] + native = identity["rank_native_libraries"] + if len(devices) != world_size or len(native) != world_size or not all(native): + raise ValueError("Complete observed device/native-library rank coverage is required") + names = {device["name"] for device in devices} + if len(names) != 1: + raise ValueError("A cost grid must record one observed GPU class") + origin = recorded_source_origin( + source_root, recorded_origin=recorded_origin, recorded_origin_sha256=recorded_origin_sha256 + ) + if identity.get("recorded_source_origin", origin) != origin: + raise ValueError("Original recorded source changed during whole-step construction") + graphs = sorted({cell["graph_batch_size"] for cell in identity["profile"]["cells"]}) + runtime = { + "runtime": identity["runtime"], + "rank_native_libraries": native, + "loaded_source_sha256": sources, + "runtime_configuration": identity["runtime_configuration"], + "model_config_sha256": identity["model_config_sha256"], + "profile": identity["profile"], + } + return { + "global_graph_batch_sizes": [g * world_size for g in graphs], + "gpu": next(iter(names)), + "gpu_count": world_size, + "gpu_snapshot_sha256": _canonical_sha256(devices), + "max_draft_len": identity["profile"]["physical_k"], + "rank_local_graph_batch_sizes": graphs, + "runtime_snapshot": _canonical_sha256(runtime), + **origin, + "topology": json.dumps(topology, sort_keys=True, separators=(",", ":")), + } + + +def _source_module(name: str, path: Path) -> ModuleType: + """Load a fixed pure-Python dependency without native package initialization.""" + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise ImportError(f"Cannot load cost tooling dependency: {path.name}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def _whole_step_payload( + data: bytes, summary: dict, *, minimum_predicted_gain: float +) -> dict[str, object]: + """Map direct complete-loop cells without adding components or missing prices.""" + if summary.get("timing_scope") != WHOLE_STEP_SCOPE: + raise ValueError( + "Graph replay observations are not schema-v2 whole-step costs; " + "use the original complete-loop whole-step producer, not relabeled timings" + ) + identity = summary["identity"] + fingerprint = identity["recorded_engine_fingerprint"] + profile = identity["profile"] + if profile["measurement_scope"] != "whole-step": + raise ValueError("Whole-step export requires the original whole-step adapter") + floors = summary["sample_floors"] + if floors["per_generation"] < 12 or floors["pooled_per_shape"] < 24: + raise ValueError("Original whole-step sample floors are 12 per generation and 24 pooled") + tables, measurements, observed = {}, [], set() + digest = hashlib.sha256(data).hexdigest() + for cell in summary["cells"]: + shape = cell["execution_shape"] + g, v = shape["graph_batch_size"], shape["verifier_budget"] + expected_layout = "uniform" if v == 0 else "packed_chain_rows" + if ( + shape["workload"] != "dspark_chain" + or shape["physical_k"] != profile["physical_k"] + or shape["attention_layout"] != expected_layout + or shape["captured_head"] != profile["record_confidence_head"] + or (g, v) in observed + ): + raise ValueError( + "Whole-step execution shape, K, layout or unique cell identity differs" + ) + samples = cell["retained_rank_max_samples_ms"] + if ( + len(samples) < floors["pooled_per_shape"] + or cell["retained_count"] != len(samples) + or any( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not 0 < value < float("inf") + for value in samples + ) + or cell["rank_max_median_ms"] != statistics.median(samples) + ): + raise ValueError("Whole-step median does not reproduce from original retained samples") + observed.add((g, v)) + table = tables.setdefault(str(g), {"token_counts": [], "step_time_ms": []}) + table["token_counts"].append(v) + table["step_time_ms"].append(cell["rank_max_median_ms"]) + measurements.append( + { + "rank_local_graph_batch_size": g, + "rank_local_verifier_budget": v, + "source_result_sha256": digest, + "step_time_ms": cell["rank_max_median_ms"], + } + ) + declared = {(cell["graph_batch_size"], cell["verifier_budget"]) for cell in profile["cells"]} + if observed != declared: + raise ValueError( + "Every declared native/compact whole-step cell requires its direct measurement" + ) + for table in tables.values(): + ordered = sorted(zip(table["token_counts"], table["step_time_ms"])) + table["token_counts"], table["step_time_ms"] = [list(values) for values in zip(*ordered)] + return { + "schema_version": 2, + "cost_tables": tables, + "engine_fingerprint": fingerprint, + "engine_fingerprint_sha256": _canonical_sha256(fingerprint), + "minimum_predicted_gain": minimum_predicted_gain, + "measurements": sorted( + measurements, + key=lambda item: ( + item["rank_local_graph_batch_size"], + item["rank_local_verifier_budget"], + ), + ), + } + + +def runtime_cost_bytes( + data: bytes, + *, + physical_k: int, + graph_batch_sizes: list[int], + minimum_predicted_gain: float = 0.01, +) -> tuple[bytes, dict[str, object]]: + """Validate reuse or serialize measured complete-loop costs into existing schema. + + A raw export summary has a different timing/identity contract. Even if a + caller toggles an eligibility flag, that cannot extend its measured scope. + Already supplied schema-v2 artifacts retain their original provenance and + values; validation is not a new GPU measurement or admission decision. + """ + utilities = _source_module( + "_dspark_cost_export_utilities", Path(__file__).with_name("verify_costs.py") + ) + payload = utilities.parse_json(data) + if not isinstance(payload, dict): + raise TypeError("Costs must contain a JSON object") + whole_step = payload.get("schema") == "verify-cost-measured-summary-v1" + if whole_step: + payload = _whole_step_payload(data, payload, minimum_predicted_gain=minimum_predicted_gain) + elif payload.get("schema") == "verify-cost-raw-v1": + raise ValueError( + "Raw observations first require exact-rank reduction from the whole-step producer" + ) + consumer_path = Path(__file__).resolve().parents[1] / "_torch/speculative/dspark_planner.py" + consumer = _source_module("_dspark_cost_export_consumer", consumer_path) + fingerprint = consumer.validate_sps_cost_table_payload( + payload, + graph_batch_sizes=graph_batch_sizes, + max_draft_len=physical_k, + validate_freshness=False, + ) + # Use the same constructor as the file loader, including its bounded exact + # candidate set and any already supplied drain metadata. No policy is added. + consumer._build_exact_sps_cost_table( + payload, max_draft_len=physical_k, validate_artifact_freshness=False + ) + output = ( + json.dumps(payload, sort_keys=True, allow_nan=False).encode() + b"\n" + if whole_step + else data + ) + return output, { + "schema": "dspark-existing-cost-structural-check-v1", + "source_artifact_sha256": hashlib.sha256(data).hexdigest(), + "runtime_cost_artifact_sha256": hashlib.sha256(output).hexdigest(), + "measured_timing_scope": WHOLE_STEP_SCOPE if whole_step else "original-supplied-artifact", + "recorded_engine_fingerprint_sha256": payload["engine_fingerprint_sha256"], + "recorded_source_head": fingerprint["source_head"], + "physical_k": physical_k, + "rank_local_graph_batch_sizes": graph_batch_sizes, + "structural_validation_passed": True, + "destination_freshness_checked": False, + "new_measurements_created": False, + "direct_whole_step_measurements_exported": whole_step, + "task_accuracy_qualified": False, + "performance_or_economic_admission_qualified": False, + "warning": ( + "Supplied predictions retain their original source and GPU identity; " + "their freshness and performance on the destination are unchecked." + ), + } + + +def check_cost_bytes( + data: bytes, *, physical_k: int, graph_batch_sizes: list[int] +) -> dict[str, object]: + """Return the existing artifact or complete-loop export's structural report.""" + return runtime_cost_bytes(data, physical_k=physical_k, graph_batch_sizes=graph_batch_sizes)[1] + + +def _write_new_bytes(path: Path, data: bytes) -> None: + """Preserve every original byte; never overwrite or follow a final symlink.""" + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0) + with os.fdopen(os.open(path, flags, 0o600), "wb") as stream: + stream.write(data) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input", type=Path, required=True) + parser.add_argument("--physical-k", type=int, required=True) + parser.add_argument("--graph-batch-sizes", type=int, nargs="+", required=True) + parser.add_argument( + "--minimum-predicted-gain", + type=float, + default=0.01, + help="Policy margin for a newly exported whole-step table only.", + ) + parser.add_argument( + "--output", type=Path, help="Optional new copy; original bytes are preserved." + ) + args = parser.parse_args(argv) + if args.input.is_symlink() or not args.input.is_file(): + raise ValueError("Input must be a regular, non-symlink cost artifact") + data = args.input.read_bytes() + output, report = runtime_cost_bytes( + data, + physical_k=args.physical_k, + graph_batch_sizes=args.graph_batch_sizes, + minimum_predicted_gain=args.minimum_predicted_gain, + ) + if args.input.read_bytes() != data: + raise ValueError("Original cost artifact changed during validation") + if args.output is not None: + _write_new_bytes(args.output, output) + print(json.dumps(report, sort_keys=True, allow_nan=False)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tensorrt_llm/tools/dspark_observations.py b/tensorrt_llm/tools/dspark_observations.py new file mode 100644 index 000000000000..d6745a4f4a20 --- /dev/null +++ b/tensorrt_llm/tools/dspark_observations.py @@ -0,0 +1,213 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU joins for offline DSpark observations; no timing or policy decisions.""" + +from __future__ import annotations + +import dataclasses +import hashlib +import json +import math +from collections import Counter +from typing import Any + +import torch + + +def token_digest(tokens: list[int]) -> str: + return hashlib.sha256(json.dumps(tokens, separators=(",", ":")).encode()).hexdigest() + + +@dataclasses.dataclass +class ProposalSnapshot: + """One genuine forward's proposals and confidence under the same sampler event. + + Host tensors must not be read until the ORIGINAL sampler has updated this + state. ``device`` ownership identifies which proposal the later forward + consumed; row order or the latest live confidence never establishes pairing. + """ + + iteration: int + state: Any + confidence_host: torch.Tensor + requests: dict[int, tuple[Any, int, int]] # original object, sampler slot, confidence row + stamp_host: torch.Tensor + draft_sequence: int + updated: bool = False + + +class NativePairs: + """Accept only uncensored, incarnation- and proposal-exact full-K pairs.""" + + def __init__(self, physical_k: int) -> None: + if type(physical_k) is not int or physical_k < 1: + raise ValueError("Positive physical K required") + self.physical_k = physical_k + self.logits: list[torch.Tensor] = [] + self.accepted: list[int] = [] + self.joins: list[dict] = [] + self._pair_keys: set[tuple[int, int, int, int]] = set() + self.stats: Counter = Counter() + + def record( + self, + *, + source: ProposalSnapshot, + request: Any, + target_iteration: int, + input_proposal: torch.Tensor, + accepted: int, + verified: int, + target_slot: int, + ) -> None: + if not source.updated: + raise RuntimeError("Original source sampler update has not completed") + if target_iteration <= source.iteration: + raise ValueError("A proposal must be verified after its producing forward") + record = source.requests.get(request.py_request_id) + if record is None or record[0] is not request: + raise ValueError("Proposal belongs to another request incarnation") + if type(accepted) is not int or not 0 <= accepted <= self.physical_k: + raise ValueError("Full-K accepted-prefix count outside physical K") + if verified != self.physical_k: + raise ValueError("Compact or partial verification censors calibration labels") + _, source_slot, confidence_row = record + if ( + source.stamp_host.device.type != "cpu" + or int(source.stamp_host[confidence_row]) != source.draft_sequence + ): + raise ValueError("Confidence row was not written by the bound producing draft") + source_tokens = source.state.host.next_draft_tokens[source_slot] + if ( + source_tokens.device.type != "cpu" + or input_proposal.device.type != "cpu" + or source_tokens.numel() != self.physical_k + or input_proposal.numel() != self.physical_k + or not torch.equal(source_tokens, input_proposal) + ): + raise ValueError("Actual verification input differs from the bound source proposal") + logits = source.confidence_host[confidence_row] + if ( + logits.device.type != "cpu" + or logits.shape != (self.physical_k,) + or not logits.is_floating_point() + or not torch.isfinite(logits).all() + ): + raise ValueError("Bound confidence must contain finite full-K logits") + key = (request.py_request_id, id(request), source.iteration, target_iteration) + if key in self._pair_keys: + raise ValueError("Repeated calibration pair") + self._pair_keys.add(key) + self.logits.append(logits.clone()) + self.accepted.append(accepted) + self.joins.append( + { + "pair_key": list(key), + "source_slot": source_slot, + "target_slot": target_slot, + "confidence_row": confidence_row, + "proposal_sha256": token_digest(source_tokens.tolist()), + "accepted_draft_prefix": accepted, + "verified_drafts": verified, + "producing_draft_sequence": source.draft_sequence, + } + ) + self.stats["recorded"] += 1 + + def payload(self) -> dict: + if not self.logits: + raise ValueError("No genuine paired full-K observations") + counts = torch.tensor(self.accepted, dtype=torch.int64).view(-1, 1) + return { + "logits": torch.stack(self.logits), + "prefix_mask": (torch.arange(self.physical_k).view(1, -1) < counts).float(), + "meta": { + "pairing": "draft_seq_ring", + "pairing_protocol": "original-input-state-object-and-proposal-v1", + "block_size": self.physical_k, + "stats": dict(self.stats), + "joins": self.joins, + "native_full_k": True, + "task_accuracy_qualified": False, + }, + } + + +def validate_responses(rows: list[dict], inputs: list[dict], output_tokens: int) -> None: + if len(rows) != len(inputs) or len({row["input_id"] for row in rows}) != len(rows): + raise ValueError("Missing, duplicated or additional original response") + expected = {row["input_id"]: row["prompt_token_ids"] for row in inputs} + for row in rows: + tokens = row["output_token_ids"] + if ( + row["input_id"] not in expected + or row["prompt_token_ids"] != expected[row["input_id"]] + or row["finish_reason"] != "length" + or len(tokens) != output_tokens + or any(type(token) is not int or token < 0 for token in tokens) + ): + raise ValueError("Original prompt, output cardinality or token contract failed") + + +def compare_runs(records: list[dict]) -> dict: + """Matched fixed-cohort E2E comparison, not task accuracy or a best-K claim.""" + arms: dict[str, list[dict]] = {"static": [], "dynamic": []} + contract = None + for record in records: + if ( + record.get("schema") != "dspark-public-run-v1" + or not record.get("run_complete") + or not record.get("model_shutdown_completed") + ): + raise ValueError("Only naturally completed, closed public runs can be compared") + mode = record["mode"] + if mode not in arms or record.get("calibration_collection"): + raise ValueError("Calibration observation is not a serving benchmark arm") + match = record["comparison_contract"] + if contract is None: + contract = match + if match != contract: + raise ValueError("Inputs, physical K, BS or numerical runtime differ") + if bool(record["resolved_confidence_enabled"]) != (mode == "dynamic") or bool( + record["runtime_confidence_enabled"] + ) != (mode == "dynamic"): + raise ValueError("Requested dynamic execution fell back to static") + routes = record["routing"] + if not routes["native_replays"] and not routes["compact_replays"]: + raise ValueError("No real generation graph replay") + if mode == "static" and routes["compact_replays"]: + raise ValueError("Static arm unexpectedly used compact verification") + if mode == "dynamic" and routes["compact_replays"] <= 0: + raise ValueError("Enabled policy never executed compact; no dynamic gain claim") + value = record["elapsed_seconds"] + if type(value) not in (int, float) or not math.isfinite(value) or value <= 0: + raise ValueError("Positive real E2E elapsed time required") + validate_responses(record["responses"], record["inputs"], contract["output_tokens"]) + arms[mode].append(record) + if not all(arms.values()) or len(arms["static"]) != len(arms["dynamic"]): + raise ValueError("Equal numbers of matched static and dynamic lifetimes required") + token_work = contract["batch_size"] * contract["output_tokens"] + prices = { + arm: token_work * len(runs) / sum(run["elapsed_seconds"] for run in runs) + for arm, runs in arms.items() + } + native_outputs = { + row["input_id"]: row["output_token_ids"] for row in arms["static"][0]["responses"] + } + equal = all( + {row["input_id"]: row["output_token_ids"] for row in run["responses"]} == native_outputs + for runs in arms.values() + for run in runs + ) + return { + "schema": "dspark-public-comparison-v1", + "comparison_contract": contract, + "timing_scope": "fixed-cohort-request-enqueue-to-final-response-including-prefill-and-drain", + "tps": prices, + "dynamic_over_static_pct": 100 * (prices["dynamic"] / prices["static"] - 1), + "exact_output_tokens_equal": equal, + "task_accuracy_qualified": False, + "default_production_gain_qualified": False, + "confidence_interval": None, + "best_static_or_dynamic_k_qualified": False, + } diff --git a/tensorrt_llm/tools/dspark_reproduction.py b/tensorrt_llm/tools/dspark_reproduction.py new file mode 100644 index 000000000000..dc318763095d --- /dev/null +++ b/tensorrt_llm/tools/dspark_reproduction.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Public native calibration collection and normal static/enabled comparison. + +Run by file path so CPU-only comparison/help does not initialize TRT-LLM/CUDA. +Each ``run`` is one independent model lifetime; order/repeats belong to caller. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import sys +from pathlib import Path + + +def _cpu_module(): + path = Path(__file__).with_name("dspark_observations.py") + name = "_dspark_public_observations" + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__, allow_abbrev=False) + sub = parser.add_subparsers(dest="command", required=True) + run = sub.add_parser("run", help="One real full-model lifetime, with natural shutdown") + run.add_argument("--mode", choices=("static", "dynamic", "collect"), required=True) + run.add_argument("--model", type=Path, required=True) + run.add_argument("--config", type=Path, required=True) + run.add_argument("--inputs", type=Path, required=True) + run.add_argument("--source-root", type=Path, required=True) + run.add_argument( + "--recorded-source-origin", + type=Path, + help="Optional recorded HEAD/full-diff JSON for this immutable checkout.", + ) + run.add_argument( + "--recorded-source-origin-sha256", + help="SHA256 of the exact optional recorded source-origin bytes.", + ) + run.add_argument("--physical-k", type=int, required=True) + run.add_argument( + "--batch-size", type=int, required=True, help="Global fixed-cohort request count" + ) + run.add_argument("--output-tokens", type=int, required=True) + run.add_argument( + "--timeout", type=float, required=True, help="Full startup+model+shutdown budget in seconds" + ) + run.add_argument("--warmup-cohorts", type=int, default=1) + run.add_argument("--cost-table", type=Path) + run.add_argument("--calibration", type=Path) + run.add_argument("--output", type=Path, required=True, help="Fresh private output directory") + compare = sub.add_parser("compare", help="Compare closed independent model lifetimes") + compare.add_argument("--runs", type=Path, nargs="+", required=True) + compare.add_argument("--output", type=Path, required=True) + args = parser.parse_args(argv) + if args.command == "run": + from tensorrt_llm._torch.speculative.dspark_reproduction import run_native + + return run_native(args) + records = [json.loads(path.read_bytes()) for path in args.runs] + report = _cpu_module().compare_runs(records) + with args.output.open("x", encoding="utf-8") as handle: + json.dump(report, handle, indent=2, allow_nan=False) + + +if __name__ == "__main__": + main() diff --git a/tensorrt_llm/tools/verify_costs.py b/tensorrt_llm/tools/verify_costs.py new file mode 100644 index 000000000000..1888a83bf41f --- /dev/null +++ b/tensorrt_llm/tools/verify_costs.py @@ -0,0 +1,445 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Workload-independent offline GPU measurements and strict raw-cost export. + +This module needs no CUDA or native TRT-LLM import to reduce existing measurements. +Its exports are not serving cost tables, accuracy certificates or admission receipts. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +import re +import statistics +from collections import defaultdict +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Callable, Protocol + +__all__ = [ + "EventLeasePool", + "canonical_sha256", + "parse_json", + "sha256_file", + "reduce_measurements", + "write_new_json", + "isolated_control_communicator", + "create_private_output_directory", +] + + +class TimingEvent(Protocol): + def record(self, stream=None) -> None: ... + def query(self) -> bool: ... + def elapsed_time(self, other: TimingEvent) -> float: ... + + +class ControlCommunicator(Protocol): + def Dup(self) -> ControlCommunicator: ... + def Free(self) -> None: ... + + +@contextmanager +def isolated_control_communicator(runtime: ControlCommunicator) -> Iterator[ControlCommunicator]: + """Duplicate collectively before workers; never modify the runtime communicator. + + The caller must finish worker shutdown and evidence closure inside this + context. Its dedicated communicator remains owned on every error path. + """ + control = runtime.Dup() + try: + yield control + finally: + control.Free() + + +class EventLeasePool: + """Keep bounded event pairs alive until completed-row readback releases them. + + The caller materializes CUDA events before capture and reads them only after + its normal completion boundary. No per-step synchronization is introduced. + """ + + def __init__(self, event_factory: Callable[[], TimingEvent], size: int = 32) -> None: + if type(size) is not int or not 2 <= size <= 256: + raise ValueError("Event pool size must be in 2..256") + self._events = [(event_factory(), event_factory()) for _ in range(size)] + self._owners: dict[int, int] = {} + self._ended: set[int] = set() + self._next_id = 0 + self._initialized = False + + def initialize(self, stream=None) -> None: + """Materialize lazy handles outside capture; caller establishes completion.""" + if self._initialized: + raise RuntimeError("Event pool already initialized") + for start, end in self._events: + start.record(stream) + end.record(stream) + self._initialized = True + + def begin(self) -> int: + if not self._initialized: + raise RuntimeError("Initialize timing events before capture") + occupied = set(self._owners.values()) + slot = next((i for i in range(len(self._events)) if i not in occupied), None) + if slot is None: + raise RuntimeError("No completed event lease available") + lease = self._next_id + self._next_id += 1 + self._owners[lease] = slot + self._events[slot][0].record() + return lease + + def end(self, lease: int) -> None: + if type(lease) is not int or lease not in self._owners: + raise RuntimeError("Unknown timing lease") + if lease in self._ended: + raise RuntimeError("Timing lease ended twice") + self._events[self._owners[lease]][1].record() + self._ended.add(lease) + + def complete(self, lease: int) -> float: + if type(lease) is not int or lease not in self._owners: + raise RuntimeError("Unknown timing lease") + if lease not in self._ended: + raise RuntimeError("Timing lease has no end boundary") + start, end = self._events[self._owners[lease]] + if not end.query(): + raise RuntimeError("Normal completion did not finish the timing events") + elapsed = float(start.elapsed_time(end)) + if not math.isfinite(elapsed) or elapsed <= 0: + raise ValueError("Measured duration must be positive and finite") + del self._owners[lease] + self._ended.remove(lease) + return elapsed + + @property + def pending(self) -> int: + return len(self._owners) + + +def canonical_sha256(value: object) -> str: + return hashlib.sha256( + json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() + ).hexdigest() + + +def parse_json(data: bytes | str) -> object: + """Reject ambiguous duplicate keys and nonstandard nonfinite JSON values.""" + + def pairs(items): + result = {} + for key, value in items: + if key in result: + raise ValueError(f"Duplicate JSON key: {key}") + result[key] = value + return result + + def constant(value): + raise ValueError(f"Nonfinite JSON value: {value}") + + return json.loads(data, object_pairs_hook=pairs, parse_constant=constant) + + +def sha256_file(path: Path) -> str: + """Hash actual artifact bytes without loading a whole native library into RAM.""" + digest = hashlib.sha256() + with path.open("rb") as stream: + while block := stream.read(1 << 20): + digest.update(block) + return digest.hexdigest() + + +def write_new_json(path: Path, value: object) -> None: + """Create a new owner-only POSIX file; never follow/overwrite an existing target. + + No existing ACL, umask or host setting is changed. POSIX permission semantics + need Linux qualification; Windows ACL behavior is not attested by this helper. + """ + payload = json.dumps(value, sort_keys=True, allow_nan=False) + "\n" + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0) + with os.fdopen(os.open(path, flags, 0o600), "w", encoding="utf-8") as stream: + stream.write(payload) + + +def create_private_output_directory(path: Path) -> None: + """Create only a new diagnostic directory, with owner-only POSIX permissions.""" + path.mkdir(mode=0o700, parents=False, exist_ok=False) + + +def reduce_measurements( + documents: list[dict], + *, + excluded_generations: int = 1, + warmup_samples: int = 2, + minimum_samples: int = 24, + minimum_samples_per_generation: int = 12, +) -> dict: + """Join exact rank/sample identities before reducing rank maxima. + + Warm-up exclusions remain in the output. Missing ranks, unequal ordinal sets, + mixed identities or scopes, and incomplete lifecycle closure are errors. + Samples are not summed across ranks or across overlapping timing scopes. + """ + for value in (excluded_generations, warmup_samples): + if type(value) is not int or value < 0: + raise ValueError("Warm-up exclusions must be nonnegative integers") + if ( + type(minimum_samples) is not int + or minimum_samples < 1 + or not documents + or type(minimum_samples_per_generation) is not int + or minimum_samples_per_generation < 1 + ): + raise ValueError("A positive sample floor and raw documents are required") + if any(not isinstance(row, dict) for row in documents): + raise ValueError("Raw measurements must be JSON objects") + first = documents[0] + world = first["world_size"] + if type(world) is not int or world < 1: + raise ValueError("Actual world size required") + identity = first["identity"] + scope = first["timing_scope"] + if not isinstance(scope, str) or not scope or not isinstance(identity, dict) or not identity: + raise ValueError("Observed identity and explicit timing scope required") + declared = identity.get("declared_shape_ids", []) + if ( + not isinstance(declared, list) + or not declared + or any( + not isinstance(key, str) or not re.fullmatch(r"[0-9a-f]{64}", key) for key in declared + ) + or len(set(declared)) != len(declared) + ): + raise ValueError("Adapter-declared finite shape identities required") + expected_generations = identity.get("generations_per_shape") + if type(expected_generations) is not int or expected_generations < 1: + raise ValueError("Declared generations per shape required") + seen_shapes = set() + declared_to_observed = {} + cells: dict[str, dict[int, dict[int, dict]]] = defaultdict(lambda: defaultdict(dict)) + shapes = {} + raw_hashes = [] + for row in documents: + if ( + row.get("schema") != "verify-cost-raw-v1" + or type(row["world_size"]) is not int + or row["world_size"] != world + or row["identity"] != identity + or row["timing_scope"] != scope + or row.get("generation_closed") is not True + or row.get("run_complete") is not True + or row.get("model_shutdown_completed") is not True + ): + raise ValueError("Raw source, timing scope, generation or shutdown closure differs") + rank, generation = row["rank"], row["generation"] + if ( + type(rank) is not int + or not 0 <= rank < world + or type(generation) is not int + or not 0 <= generation < expected_generations + ): + raise ValueError("Invalid rank/generation identity") + if not isinstance(row["execution_shape"], dict) or not row["execution_shape"]: + raise ValueError("Explicit execution shape/layout required") + key = canonical_sha256(row["execution_shape"]) + if row.get("declared_shape_id") not in declared: + raise ValueError("Measurement is not part of the declared workload grid") + seen_shapes.add(row["declared_shape_id"]) + observed = declared_to_observed.setdefault(row["declared_shape_id"], key) + if observed != key: + raise ValueError("A declared shape changed its observed execution layout") + shapes[key] = row["execution_shape"] + if rank in cells[key][generation]: + raise ValueError("Duplicate cell/generation/rank") + cells[key][generation][rank] = row + raw_hashes.append(canonical_sha256(row)) + if seen_shapes != set(declared): + raise ValueError("A whole declared workload shape is missing") + results = [] + rank_ordinals = [dict() for _ in range(world)] + for key, generations in sorted(cells.items()): + if sorted(generations) != list(range(expected_generations)): + raise ValueError("Declared per-cell generation coverage is incomplete") + retained, excluded, joined_rows = [], [], [] + previous_maximum = -1 + for generation, ranks in sorted(generations.items()): + if set(ranks) != set(range(world)): + raise ValueError("Missing rank in a generation") + samples = [] + for rank in range(world): + rank_samples = {} + if not isinstance(ranks[rank]["samples"], list) or not ranks[rank]["samples"]: + raise ValueError("A generation has no raw samples") + for sample in ranks[rank]["samples"]: + ordinal, iteration = sample["sample_index"], sample["executor_iteration"] + duration = sample["elapsed_ms"] + if ( + type(ordinal) is not int + or ordinal < 0 + or ordinal in rank_samples + or type(iteration) is not int + or iteration < 0 + or type(duration) not in (int, float) + or not math.isfinite(duration) + or duration <= 0 + or type(sample["eligible"]) is not bool + ): + raise ValueError("Invalid or duplicated raw sample") + rank_samples[ordinal] = sample + if ordinal in rank_ordinals[rank]: + raise ValueError("A replay ordinal was reused in the run") + rank_ordinals[rank][ordinal] = iteration + samples.append(rank_samples) + ordinals = set(samples[0]) + if any(set(rank_samples) != ordinals for rank_samples in samples): + raise ValueError("Ranks have different replay ordinal sets") + if min(ordinals) <= previous_maximum: + raise ValueError("Replay ordinals reset between generations") + previous_maximum = max(ordinals) + eligible_index = 0 + generation_retained = 0 + for ordinal in sorted(ordinals): + joined = [samples[rank][ordinal] for rank in range(world)] + if len({sample["executor_iteration"] for sample in joined}) != 1: + raise ValueError("Replay ordinal joins different executor iterations") + duration = max(sample["elapsed_ms"] for sample in joined) + reason = None + if not all(sample["eligible"] for sample in joined): + reason = "non_plateau_or_incomplete_update" + elif generation < excluded_generations: + reason = "warmup_generation" + else: + if eligible_index < warmup_samples: + reason = "initial_plateau_sample" + eligible_index += 1 + detail = { + "generation": generation, + "sample_index": ordinal, + "executor_iteration": joined[0]["executor_iteration"], + "rank_elapsed_ms": [sample["elapsed_ms"] for sample in joined], + "rank_max_elapsed_ms": duration, + "excluded_reason": reason, + } + joined_rows.append(detail) + (excluded if reason else retained).append(duration) + if reason is None: + generation_retained += 1 + if ( + generation >= excluded_generations + and generation_retained < minimum_samples_per_generation + ): + raise ValueError("Insufficient retained exact samples in a generation") + if len(retained) < minimum_samples: + raise ValueError("Insufficient exact retained samples for a cell") + results.append( + { + "execution_shape": shapes[key], + "retained_count": len(retained), + "rank_max_median_ms": statistics.median(retained), + "retained_rank_max_samples_ms": retained, + "joins": joined_rows, + "excluded_count": len(excluded), + } + ) + for ordinals in rank_ordinals: + ordered = sorted(ordinals) + if ordered != list(range(len(ordered))): + raise ValueError("Whole-run replay ordinal coverage is incomplete") + iterations = [ordinals[ordinal] for ordinal in ordered] + if any(right <= left for left, right in zip(iterations, iterations[1:])): + raise ValueError("Executor iterations do not advance with replay ordinals") + return { + "schema": "verify-cost-measured-summary-v1", + "identity": identity, + "timing_scope": scope, + "cells": results, + "raw_document_sha256": raw_hashes, + "exclusions": {"generations": excluded_generations, "initial_samples": warmup_samples}, + "sample_floors": { + "per_generation": minimum_samples_per_generation, + "pooled_per_shape": minimum_samples, + }, + "production_sps_table_eligible": False, + "production_admission_qualified": False, + "policy_overhead_measured": False, + "task_accuracy_qualified": False, + } + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest="command", required=True) + collect = commands.add_parser("profile", help="Collect authentic DSpark graph replays") + collect.add_argument("--model", type=Path, required=True) + collect.add_argument("--config", type=Path, required=True) + collect.add_argument("--grid", type=Path, required=True) + collect.add_argument("--physical-k", type=int, required=True) + collect.add_argument( + "--cost-scope", choices=("graph-replay", "whole-step"), default="graph-replay" + ) + collect.add_argument( + "--source-root", + type=Path, + help="Original public checkout used for recorded whole-step metadata.", + ) + collect.add_argument("--inputs", type=Path, required=True, help="JSONL prompt_token_ids") + collect.add_argument( + "--recorded-source-origin", + type=Path, + help="Optional recorded HEAD/full-diff JSON for this immutable checkout.", + ) + collect.add_argument( + "--recorded-source-origin-sha256", + help="SHA256 of the exact optional recorded source-origin bytes.", + ) + collect.add_argument("--output", type=Path, required=True) + collect.add_argument("--generations", type=int, default=3) + collect.add_argument("--output-tokens", type=int, default=768) + collect.add_argument("--timeout", type=float, default=12180.0) + export = commands.add_parser("export", help="Reduce raw observations; never grant admission") + export.add_argument("--raw", type=Path, nargs="+", required=True) + export.add_argument("--output", type=Path, required=True) + export.add_argument("--exclude-generations", type=int, default=1) + export.add_argument("--warmup-samples", type=int, default=2) + export.add_argument("--minimum-samples", type=int, default=24) + export.add_argument("--minimum-generation-samples", type=int, default=12) + args = parser.parse_args(argv) + if ( + args.command == "profile" + and args.cost_scope != "whole-step" + and ( + args.recorded_source_origin is not None + or args.recorded_source_origin_sha256 is not None + ) + ): + parser.error("Recorded source-origin input applies only to --cost-scope whole-step") + if args.command == "profile": + from tensorrt_llm._torch.speculative.dspark_cost_profiler import run_profile + + run_profile(args) + else: + raw_sources = [(path, path.read_bytes()) for path in args.raw] + documents = [parse_json(raw) for _, raw in raw_sources] + summary = reduce_measurements( + documents, + excluded_generations=args.exclude_generations, + warmup_samples=args.warmup_samples, + minimum_samples=args.minimum_samples, + minimum_samples_per_generation=args.minimum_generation_samples, + ) + summary["raw_file_sha256"] = { + str(path): hashlib.sha256(raw).hexdigest() for path, raw in raw_sources + } + write_new_json(args.output, summary) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index b96235b86327..cd6bbd97f36b 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1679,6 +1679,16 @@ "kind": "categorical", "path": "speculative_config.eagle3_model_arch" }, + { + "capture_policy": "bool", + "kind": "value", + "path": "speculative_config.enable_confidence_scheduling" + }, + { + "capture_policy": "bool", + "kind": "value", + "path": "speculative_config.enable_fused_confidence_scheduler" + }, { "capture_policy": "bool", "kind": "value", @@ -1838,6 +1848,11 @@ "kind": "value", "path": "speculative_config.use_relaxed_acceptance_for_thinking" }, + { + "capture_policy": "bool", + "kind": "value", + "path": "speculative_config.validate_confidence_artifact_freshness" + }, { "capture_policy": "int", "kind": "value", diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 90f9bb55eb51..13b447516324 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -89,6 +89,7 @@ l0_cpu: - unittest/_torch/multimodal - unittest/_torch/ray_orchestrator/single_gpu/test_cache_transceiver_comm.py - unittest/_torch/speculative/hw_agnostic + - unittest/_torch/speculative/test_device_window_sampling_layout.py # l0_h100 selects with -m "not cpu_only", so this entry is what runs it. - unittest/_torch/speculative/test_spec_slot_pool_sizing.py - unittest/_torch/test_hf_quant_config.py @@ -150,6 +151,7 @@ l0_cpu: - unittest/executor/test_stop_word_side_output_trim.py - unittest/executor/test_ray_stub.py - unittest/executor/test_generation_request_array_prompt.py + - unittest/executor/test_dspark_reproduction_factory.py - unittest/grpc/openengine/test_coordinator.py - unittest/grpc/openengine/test_frontend_lifecycle.py - unittest/grpc/openengine/test_control.py @@ -192,6 +194,7 @@ l0_cpu: - unittest/llmapi/test_whisper_suppress_tokens_processor.py - unittest/llmapi/test_kv_cache_dtype_override.py - unittest/llmapi/test_llm_args.py + - unittest/llmapi/test_dspark_workload_admission_config.py - unittest/llmapi/test_mx_args.py - unittest/llmapi/test_custom_tokenizer_aliases.py - unittest/llmapi/test_llm_quant.py diff --git a/tests/unittest/_torch/executor/engine/test_metadata.py b/tests/unittest/_torch/executor/engine/test_metadata.py index 801d01b70f2f..392afda5f521 100644 --- a/tests/unittest/_torch/executor/engine/test_metadata.py +++ b/tests/unittest/_torch/executor/engine/test_metadata.py @@ -95,3 +95,47 @@ def test_build_attention_metadata_forwards_shared_and_cache_inputs() -> None: assert metadata.draft_kv_cache_manager is draft_kv_cache_manager assert metadata.enable_context_mla_with_cached_kv assert metadata.num_heads_per_kv == 4 + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_build_attention_metadata_passes_ragged_opt_in_to_supported_backend(enabled) -> None: + class RaggedMetadata(_AttentionMetadata): + enable_ragged_verification = False + + model_config = SimpleNamespace( + pretrained_config=SimpleNamespace(), + sparse_attention_config=None, + enable_flash_mla=False, + ) + metadata = build_attention_metadata( + model_config, + max_batch_size=4, + max_num_tokens=16, + max_beam_width=1, + attention_backend=SimpleNamespace(Metadata=RaggedMetadata), + attention_runtime_features=AttentionRuntimeFeatures(), + mapping=object(), + cache_indirection=None, + enable_ragged_verification=enabled, + ) + assert metadata.enable_ragged_verification is enabled + + +def test_build_attention_metadata_does_not_pass_ragged_option_to_other_backend() -> None: + model_config = SimpleNamespace( + pretrained_config=SimpleNamespace(), + sparse_attention_config=None, + enable_flash_mla=False, + ) + metadata = build_attention_metadata( + model_config, + max_batch_size=4, + max_num_tokens=16, + max_beam_width=1, + attention_backend=_AttentionBackend, + attention_runtime_features=AttentionRuntimeFeatures(), + mapping=object(), + cache_indirection=None, + enable_ragged_verification=True, + ) + assert not hasattr(metadata, "enable_ragged_verification") diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py index c6509d5c503e..9cbdbaddfe50 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py @@ -8,6 +8,7 @@ import tensorrt_llm import tensorrt_llm.bindings +from tensorrt_llm._torch.pyexecutor._util import _get_num_reserved_index_slots from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2 as kv_cache_v2_module from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState, SamplingConfig @@ -16,6 +17,27 @@ CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType +@pytest.mark.parametrize( + ("spec_config", "estimating_kv_cache", "expected"), + [ + (None, False, 1), + (SimpleNamespace(enable_confidence_scheduling=False), False, 1), + (SimpleNamespace(enable_confidence_scheduling=True), False, 2), + (SimpleNamespace(enable_confidence_scheduling=True), True, 0), + ], + ids=[ + "no-speculation", + "static-dspark", + "confidence-dspark", + "confidence-estimator", + ], +) +def test_index_mapper_reservation_tracks_confidence_scheduling( + spec_config: SimpleNamespace | None, estimating_kv_cache: bool, expected: int +) -> None: + assert _get_num_reserved_index_slots(spec_config, estimating_kv_cache) == expected + + def _manager( *, is_draft: bool, diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index a2ed4a1b20eb..b76cd5c3b1dc 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -828,6 +828,7 @@ def test_graph_key_forwards_promoted_context_ids(self) -> None: runner.config = SimpleNamespace() runner._resolve_sample_type.return_value = SampleType.FULL runner._get_seq_len_mode.return_value = True + runner._ragged_verify_bucket.return_value = None request = _make_request_stub(7) batch = ScheduledRequests() batch.generation_requests = [request] @@ -856,6 +857,7 @@ def test_graph_key_aggregates_encoder_tokens(self) -> None: runner._resolve_sample_type.return_value = SampleType.FULL runner.max_beam_width = 1 runner._get_seq_len_mode.return_value = False + runner._ragged_verify_bucket.return_value = None context = _make_request_stub(1) context.encoder_output_len = 7 context.py_skip_cross_kv_projection = False @@ -880,6 +882,7 @@ def test_graph_key_rejects_nonuniform_context_query_lengths(self) -> None: runner.config = SimpleNamespace() runner._resolve_sample_type.return_value = SampleType.FULL runner._get_seq_len_mode.return_value = False + runner._ragged_verify_bucket.return_value = None first_context = _make_request_stub(1) first_context.encoder_output_len = 7 first_context.py_skip_cross_kv_projection = False @@ -927,6 +930,7 @@ def test_graph_key_includes_peft_cache_dtype(self) -> None: runner.config = SimpleNamespace() runner._resolve_sample_type.return_value = SampleType.FULL runner._get_seq_len_mode.return_value = False + runner._ragged_verify_bucket.return_value = None request = _make_request_stub(7) batch = ScheduledRequests() batch.generation_requests = [request] @@ -945,6 +949,7 @@ def test_graph_key_includes_peft_cache_dtype(self) -> None: def test_graph_dtype_change_falls_back_to_eager(self) -> None: runner = Mock() runner.enabled = True + runner._dspark_confidence_enabled = False runner.config = SimpleNamespace( enable_attention_dp=False, use_mrope=False, @@ -985,6 +990,7 @@ def test_graph_key_includes_lora_variant(self) -> None: runner.config = SimpleNamespace() runner._resolve_sample_type.return_value = SampleType.FULL runner._get_seq_len_mode.return_value = False + runner._ragged_verify_bucket.return_value = None request = _make_request_stub(7) batch = ScheduledRequests() batch.generation_requests = [request] @@ -1027,6 +1033,7 @@ def test_lora_graph_variant_selection(self) -> None: def test_graph_lookup_forwards_promoted_context_ids(self) -> None: runner = Mock() runner.enabled = True + runner._dspark_confidence_enabled = False runner.config = SimpleNamespace( enable_attention_dp=False, use_mrope=False, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_cuda_graph_dummy_identity.py b/tests/unittest/_torch/speculative/hw_agnostic/test_cuda_graph_dummy_identity.py new file mode 100644 index 000000000000..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..6b8a670def45 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph_runtime.py @@ -0,0 +1,385 @@ +# 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 +import torch + +from tensorrt_llm._torch.pyexecutor import cuda_graph_runner as graph_module +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( + CUDA_GRAPH_DUMMY_REQUEST_ID, + ADPShapeAgreement, + CUDAGraphRunner, + DSparkCUDAGraphRunner, + KeyType, + cuda_graph_dummy_request_id, +) +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.sampler.sampler_common import SampleType + + +class _Batch: + def __init__(self, generation_requests): + self.generation_requests = list(generation_requests) + self.context_requests = [] + self.encoder_requests = [] + self.can_run_cuda_graph = True + + @property + def batch_size(self): + return len(self.generation_requests) + + @property + def num_context_requests(self): + return 0 + + @property + def num_generation_requests(self): + return len(self.generation_requests) + + +def _runner_stub(runner_type=DSparkCUDAGraphRunner): + runner = runner_type.__new__(runner_type) + runner.clear = lambda: None + return runner + + +def _agreement(batch, *, iteration=9, peer_sizes=(3, 4), reuse=True): + return ADPShapeAgreement( + iteration=iteration, + batch_identity=id(batch), + local_batch_size=batch.batch_size, + peer_batch_sizes=peer_sizes, + all_can_graph=True, + widest_batch_size=max(peer_sizes), + graph_batch_size=4, + draft_len=5, + padding_ready=True, + ragged_bucket=16, + reuse_graph_shape=reuse, + ) + + +def test_padding_dummy_ids_are_disjoint(): + ids = { + cuda_graph_dummy_request_id(draft_len, variant=variant, max_draft_len=5) + for variant in (0, 1) + for draft_len in range(6) + } + + assert len(ids) == 12 + assert ( + cuda_graph_dummy_request_id(5, variant=0, max_draft_len=5) + == CUDA_GRAPH_DUMMY_REQUEST_ID - 5 + ) + with pytest.raises(ValueError, match="invalid CUDA padding dummy identity"): + cuda_graph_dummy_request_id(5, variant=2, max_draft_len=5) + + +def test_shape_agreement_is_bound_to_iteration_batch_and_phase(): + batch = _Batch([object()] * 3) + agreement = _agreement(batch) + + assert agreement.matches(batch, 9, padded=False) + assert not agreement.matches(batch, 8, padded=False) + assert not agreement.matches(_Batch([object()] * 3), 9, padded=False) + + batch.generation_requests.append(object()) + assert agreement.matches(batch, 9, padded=True) + assert not agreement.matches(batch, 9, padded=False) + + +def test_ragged_graph_key_requires_every_row_and_the_agreed_bucket(): + runner = _runner_stub() + runner.agreed_ragged_bucket = 22 + batch = _Batch( + [ + SimpleNamespace(py_verify_len=3), + SimpleNamespace(py_verify_len=5), + ] + ) + + assert runner._ragged_verify_bucket(batch) == 22 + batch.generation_requests[1].py_verify_len = None + assert runner._ragged_verify_bucket(batch) is None + + +def test_native_runner_does_not_inspect_ragged_request_windows(): + runner = _runner_stub(CUDAGraphRunner) + assert runner._ragged_verify_bucket(object()) is None + + +def test_ragged_graph_key_controls_the_submitted_token_count(): + runner = _runner_stub() + + assert ( + runner._get_num_tokens_for_key( + KeyType( + batch_size=4, + draft_len=5, + is_first_draft=False, + ragged_verify_bucket=14, + ) + ) + == 14 + ) + with pytest.raises(ValueError, match="mixed context"): + runner._get_num_tokens_for_key( + KeyType( + batch_size=4, + draft_len=5, + is_first_draft=False, + num_contexts=1, + ragged_verify_bucket=14, + ) + ) + + +@pytest.mark.parametrize( + ("padded_size", "num_requests", "expected"), + [(4, 3, True), (4, 4, True), (8, 3, False), (0, 3, False), (4, 0, False)], +) +def test_will_pad_to_matches_the_size_guards(padded_size, num_requests, expected): + runner = _runner_stub() + runner.supported_batch_sizes = [1, 2, 4] + runner.config = SimpleNamespace(batch_size=4) + + assert runner.will_pad_to(padded_size, num_requests) is expected + + +def test_zero_real_padding_uses_distinct_low_and_high_dummies(): + runner = _runner_stub() + runner.enabled = True + runner.padding_enabled = True + runner.max_supported_batch_size = 4 + runner.config = SimpleNamespace( + enable_attention_dp=False, + mapping=SimpleNamespace(tp_size=1), + batch_size=4, + ) + runner.adp_shape_agreement = None + runner.ragged_pad_verify_len = 2 + runner.ragged_zero_real_high_rows = 2 + runner.spec_config = SimpleNamespace(enable_confidence_scheduling=True, max_draft_len=5) + runner._can_run_cuda_graph_batch = lambda _batch: True + runner._round_up_batch_size_with_draft_len = lambda *_args: 4 + + scheduled_dummy = SimpleNamespace(is_attention_dp_dummy=True, py_verify_len=3) + low = SimpleNamespace(py_verify_len=None) + high = SimpleNamespace(py_verify_len=None) + runner._get_or_create_padding_dummy = ( + lambda _resource, _draft, variant=0: low if variant == 0 else high + ) + batch = _Batch([scheduled_dummy]) + + added = runner._get_padded_batch(batch, SimpleNamespace(), runtime_draft_len=5) + + assert added == 3 + assert batch.generation_requests == [scheduled_dummy, high, low, low] + assert [row.py_verify_len for row in batch.generation_requests] == [3, 3, 2, 2] + assert sum(1 + row.py_verify_len for row in batch.generation_requests) == 14 + + +def test_releasing_padding_dummies_invalidates_the_cached_agreement(): + runner = _runner_stub() + low = object() + high = object() + runner.padding_dummy_requests = {5: low} + runner.secondary_padding_dummy_requests = {5: high} + runner.adp_shape_agreement = object() + freed = [] + runner._padding_dummy_managers = lambda _resource_manager: [ + SimpleNamespace(free_resources=freed.append) + ] + + assert runner.release_padding_dummy(SimpleNamespace(), 5) + assert freed == [low, high] + assert runner.adp_shape_agreement is None + + +@pytest.mark.parametrize("confidence_enabled", [False, True]) +def test_decoder_selects_confidence_graph_subclass_only_when_enabled( + monkeypatch, confidence_enabled +): + monkeypatch.setattr(graph_module, "get_memory_buffers", lambda: None) + config = SimpleNamespace( + use_cuda_graph=False, + cuda_graph_padding_enabled=True, + cuda_graph_batch_sizes=[1, 4], + max_cuda_graph_batch_size=4, + max_beam_width=1, + spec_config=SimpleNamespace(enable_confidence_scheduling=confidence_enabled), + sparse_attention_config=None, + is_encoder_decoder=False, + enable_encoder_decoder_mixed_cuda_graph=False, + enable_in_graph_sampling=False, + cuda_graph_mem_pool=None, + dynamic_draft_len_mapping=None, + ) + registered = [] + decoder = SimpleNamespace( + _dspark_confidence_enabled=confidence_enabled, + _cuda_graph_runner_config=lambda: config, + model=SimpleNamespace(register_cuda_graph_pre_replay_hooks=registered.append), + ) + + runner = DecoderRunner._initialize_cuda_graph_runner(decoder) + runner.clear = lambda: None + + expected_type = DSparkCUDAGraphRunner if confidence_enabled else CUDAGraphRunner + assert type(runner) is expected_type + assert registered == [runner] + for name in ( + "agreed_ragged_bucket", + "adp_shape_agreement", + "ragged_pad_verify_len", + "ragged_zero_real_high_rows", + "secondary_padding_dummy_requests", + ): + assert hasattr(runner, name) is confidence_enabled + + +@pytest.mark.parametrize("runner_type", [CUDAGraphRunner, DSparkCUDAGraphRunner]) +def test_reset_padding_dummies_clears_only_runner_owned_state(runner_type): + runner = _runner_stub(runner_type) + runner.padding_dummy_requests = {5: object()} + if runner_type is DSparkCUDAGraphRunner: + runner.secondary_padding_dummy_requests = {5: object()} + runner.adp_shape_agreement = object() + runner.ragged_zero_real_high_rows = 2 + + runner.reset_padding_dummies() + + assert runner.padding_dummy_requests == {} + if runner_type is DSparkCUDAGraphRunner: + assert runner.secondary_padding_dummy_requests == {} + assert runner.adp_shape_agreement is None + assert runner.ragged_zero_real_high_rows == 0 + else: + assert not hasattr(runner, "secondary_padding_dummy_requests") + assert not hasattr(runner, "adp_shape_agreement") + + +def test_cached_shape_agreement_keeps_sampling_tier_collective(): + runner = _runner_stub() + batch = _Batch([SimpleNamespace(py_verify_len=3, py_draft_tokens=[0] * 5)] * 4) + runner.adp_shape_agreement = _agreement(batch, peer_sizes=(4, 4)) + runner.agreed_ragged_bucket = 16 + runner.enabled = True + runner.enable_in_graph_sampling = True + exchanges = [] + + def allgather(payload): + exchanges.append(payload) + peer = list(payload) + peer[-1] = SampleType.FULL.value + return torch.tensor([payload, peer]) + + runner.config = SimpleNamespace( + enable_attention_dp=True, + mapping=SimpleNamespace(tp_size=2), + dist=SimpleNamespace(tp_allgather_int64=allgather), + ) + runner._resolve_sample_type = lambda *_args: SampleType.FAST + + assert runner._agree_graph_shape(batch, True, frozenset()) == (True, SampleType.FULL) + assert exchanges == [[True, 4, 16, 5, SampleType.FAST.value]] + runner.enable_in_graph_sampling = False + exchanges.clear() + assert runner._agree_graph_shape(batch, True, frozenset()) == (True, None) + assert exchanges == [] + + +def test_cached_shape_agreement_rejects_a_changed_ragged_key(): + runner = _runner_stub() + batch = _Batch([SimpleNamespace(py_verify_len=3, py_draft_tokens=[0] * 5)] * 4) + runner.adp_shape_agreement = _agreement(batch, peer_sizes=(4, 4)) + runner.agreed_ragged_bucket = 15 + runner.enable_in_graph_sampling = False + + with pytest.raises(RuntimeError, match="no longer matches"): + runner._agree_graph_shape(batch, True, frozenset()) + + +def _tier_lookup_runner(batch, *, confidence_enabled, in_graph_sampling, peer_tier): + runner = _runner_stub() + payloads = [] + + def gather(payload): + payloads.append(tuple(payload)) + peer = list(payload) + peer[-1] = peer_tier.value + return np.asarray([payload, peer], dtype=np.int64) + + runner.config = SimpleNamespace( + enable_attention_dp=True, + mapping=SimpleNamespace(tp_size=2), + dist=SimpleNamespace(tp_allgather_int64=gather), + use_mrope=False, + enable_in_graph_sampling=in_graph_sampling, + ) + runner.enabled = True + runner.enable_in_graph_sampling = in_graph_sampling + runner.is_encoder_decoder = False + runner.enable_encoder_decoder_mixed_cuda_graph = False + runner._dspark_confidence_enabled = confidence_enabled + runner._dspark_trims_submitted_tokens = confidence_enabled + runner.adp_shape_agreement = None + runner.agreed_ragged_bucket = 16 if confidence_enabled else None + runner.sparse_config = None + runner._capture_sample_type = None + runner._sample_type_resolver = lambda _batch, _override: SampleType.FAST + runner._capture_allowed = False + runner.graph_metadata = { + KeyType( + batch_size=4, + draft_len=5, + is_first_draft=False, + sample_type=tier, + ragged_verify_bucket=runner.agreed_ragged_bucket, + ): {"attn_metadata": tier, "spec_metadata": None} + for tier in (SampleType.FULL, SampleType.FAST) + } + return runner, payloads + + +@pytest.mark.parametrize("confidence_enabled", [False, True]) +@pytest.mark.parametrize("peer_tier", [SampleType.FAST, SampleType.FULL]) +def test_graph_lookup_agrees_sampling_tier_without_aliasing_ragged_bucket( + confidence_enabled, peer_tier +): + batch = _Batch([SimpleNamespace(py_draft_tokens=[1] * 5, py_verify_len=3) for _ in range(4)]) + runner, payloads = _tier_lookup_runner( + batch, confidence_enabled=confidence_enabled, in_graph_sampling=True, peer_tier=peer_tier + ) + + metadata, _, key = runner.maybe_get_cuda_graph(batch, True, SimpleNamespace()) + + assert metadata is peer_tier + assert key.sample_type is peer_tier + assert payloads == [ + (True, 4, 16, 5, SampleType.FAST.value) + if confidence_enabled + else (True, 4, SampleType.FAST.value) + ] + + +@pytest.mark.parametrize("in_graph_sampling", [False, True]) +def test_cached_shape_does_not_replace_in_graph_sampling_tier_agreement(in_graph_sampling): + batch = _Batch([SimpleNamespace(py_draft_tokens=[1] * 5, py_verify_len=3) for _ in range(4)]) + runner, payloads = _tier_lookup_runner( + batch, + confidence_enabled=True, + in_graph_sampling=in_graph_sampling, + peer_tier=SampleType.FAST, + ) + runner.adp_shape_agreement = _agreement(batch, peer_sizes=(4, 4)) + + metadata, _, key = runner.maybe_get_cuda_graph(batch, True, SimpleNamespace()) + + assert metadata is (SampleType.FAST if in_graph_sampling else SampleType.FULL) + assert key.sample_type is metadata + assert len(payloads) == int(in_graph_sampling) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_dummy_slots_cpu.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_dummy_slots_cpu.py new file mode 100644 index 000000000000..7b5fecb3d949 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_dummy_slots_cpu.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Host slot-ownership regression with CPU tensors in place of CUDA allocations.""" + +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm._torch.speculative.dspark as dspark +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import cuda_graph_dummy_request_id +from tensorrt_llm._torch.speculative.dspark import DSparkSpecMetadata, DSv4DSparkWorker + +pytestmark = pytest.mark.cpu_only + + +@pytest.mark.parametrize("physical_k", range(1, 9)) +def test_all_padding_families_preserve_real_slot_ownership(monkeypatch, physical_k): + # Exercise the real initialization and host prepare methods. This deliberately + # does not validate device copies, graph replay, or target KV-cache contents. + original_zeros = torch.zeros + original_empty = torch.empty + original_full = torch.full + + def cpu_zeros(*args, **kwargs): + kwargs["device"] = "cpu" + return original_zeros(*args, **kwargs) + + def cpu_empty(*args, **kwargs): + kwargs["device"] = "cpu" + return original_empty(*args, **kwargs) + + def cpu_full(*args, **kwargs): + kwargs["device"] = "cpu" + return original_full(*args, **kwargs) + + monkeypatch.setattr(torch, "zeros", cpu_zeros) + monkeypatch.setattr(torch, "empty", cpu_empty) + monkeypatch.setattr(torch, "full", cpu_full) + monkeypatch.setattr(dspark, "prefer_pinned", lambda: False) + worker = object.__new__(DSv4DSparkWorker) + worker.spec_config = SimpleNamespace(max_draft_len=physical_k) + worker._win_inited = False + worker._attention_warmup_attempted = True + worker.return_confidence = False + worker._confidence_logits = None + worker._mixed_confidence_tracking = False + worker._confidence_incarnations = {} + draft = SimpleNamespace( + block_size=physical_k, + num_stages=1, + _attn_params={"window_size": 4, "head_dim": 2}, + ) + DSv4DSparkWorker._lazy_init(worker, draft, SimpleNamespace(max_num_requests=4)) + meta = SimpleNamespace( + _dspark_worker=worker, batch_indices_cuda=torch.empty(4, dtype=torch.int32) + ) + for variant in (0, 1): + for runtime_k in range(physical_k + 1): + dummy = cuda_graph_dummy_request_id( + runtime_k, variant=variant, max_draft_len=physical_k + ) + # Include ADP idle, the secondary zero-real family, and a live row. + meta.request_ids = [1000, 0, dummy] + meta.num_generations = 3 + DSparkSpecMetadata.prepare(meta) + assert worker._req_to_slot == {1000: 0} + assert list(worker._free_slots) == [1, 2, 3] + assert worker._batch_to_slot[:3].tolist() == [0, 4, 4] + # Removing a real request releases exactly its slot, never a dummy's. + meta.request_ids = [dummy, 0] + meta.num_generations = 2 + DSparkSpecMetadata.prepare(meta) + assert worker._req_to_slot == {} + assert sorted(worker._free_slots) == [0, 1, 2, 3] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_embedded_confidence_loading.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_embedded_confidence_loading.py new file mode 100644 index 000000000000..01881356f6e9 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_embedded_confidence_loading.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Embedded-wrapper validation through the real shared parameter walk.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch.models import modeling_dspark, modeling_dspark_v41 +from tensorrt_llm._torch.models.modeling_speculative import DSparkConfidenceHead + +pytestmark = pytest.mark.cpu_only + + +def _wrapper(monkeypatch, v41): + module = modeling_dspark_v41 if v41 else modeling_dspark + cls = module.DSv41DSparkForCausalLM if v41 else module.DSv4DSparkForCausalLM + wrapper = object.__new__(cls) + nn.Module.__init__(wrapper) + stage = nn.Module() + stage.main_proj = nn.Linear(2, 2, bias=False) + stage.main_norm = nn.RMSNorm(2) + stage.norm = nn.RMSNorm(2) + stage.markov_head = None + stage.confidence_head = DSparkConfidenceHead(hidden_size=2, block_size=2) + model = nn.Module() + model.mtp_layers = nn.ModuleList([stage]) + model.config = SimpleNamespace( + q_lora_rank=1, + num_attention_heads=1, + num_hidden_layers=0, + num_nextn_predict_layers=1, + kv_lora_rank=1, + ) + model.model_config = SimpleNamespace( + mapping=SimpleNamespace( + tp_rank=0, + tp_size=1, + cp_rank=0, + cp_size=1, + enable_attention_dp=False, + ) + ) + model.post_load_weights = Mock() + model.cache_attn_weights_from_state_dict = Mock() + wrapper.dspark_model = model + wrapper.config = model.config + wrapper.num_stages = 1 + # Remapping has separate checkpoint-layout tests. This test starts with + # model-named tensors and does not mock the loader or its parameter copies. + remapper = "remap_dspark_v41_draft_keys" if v41 else "remap_dspark_draft_keys" + monkeypatch.setattr(module, remapper, lambda weights, *args, **kwargs: weights) + weights = {name: torch.full_like(param, 2.0) for name, param in model.named_parameters()} + return wrapper, weights + + +@pytest.mark.parametrize("v41", [False, True]) +@pytest.mark.parametrize("bias", [False, True]) +def test_embedded_confidence_load_validates_and_copies_real_parameters(monkeypatch, v41, bias): + wrapper, weights = _wrapper(monkeypatch, v41) + head = wrapper.dspark_model.mtp_layers[-1].confidence_head + prefix = "mtp_layers.0.confidence_head." + if bias: + weights[prefix + "proj.bias"] = torch.tensor([3.0]) + weight_pointer = head.proj.weight.data_ptr() + rng_before = torch.get_rng_state() + wrapper.load_weights(weights) + assert head.proj.weight.data_ptr() == weight_pointer + assert torch.equal(torch.get_rng_state(), rng_before) + assert torch.equal(head.proj.weight, weights[prefix + "proj.weight"]) + assert (head.proj.bias is not None) is bias + if bias: + assert torch.equal(head.proj.bias, weights[prefix + "proj.bias"]) + wrapper.dspark_model.post_load_weights.assert_called_once() + + +@pytest.mark.parametrize("v41", [False, True]) +@pytest.mark.parametrize("fault", ["unknown", "missing_weight", "broadcast_weight"]) +def test_embedded_confidence_rejects_incomplete_or_unknown_tensors(monkeypatch, v41, fault): + wrapper, weights = _wrapper(monkeypatch, v41) + prefix = "mtp_layers.0.confidence_head." + if fault == "unknown": + weights[prefix + "bias"] = torch.tensor([99.0]) + elif fault == "missing_weight": + weights.pop(prefix + "proj.weight") + weights[prefix + "proj.bias"] = torch.tensor([3.0]) + else: + weights[prefix + "proj.weight"] = torch.tensor([[2.0]]) + with pytest.raises(ValueError, match="confidence head checkpoint"): + wrapper.load_weights(weights) + wrapper.dspark_model.post_load_weights.assert_not_called() + + +def test_v41_entirely_absent_confidence_head_remains_supported(monkeypatch): + wrapper, weights = _wrapper(monkeypatch, True) + weights.pop("mtp_layers.0.confidence_head.proj.weight") + wrapper.load_weights(weights) + assert wrapper.dspark_model.mtp_layers[-1].confidence_head is None + wrapper.dspark_model.post_load_weights.assert_called_once() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_executor_policy.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_executor_policy.py new file mode 100644 index 000000000000..f7d80b2e0110 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_executor_policy.py @@ -0,0 +1,326 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only tests for DSpark policy agreement in the executor.""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ADPShapeAgreement +from tensorrt_llm._torch.pyexecutor.py_executor import ( + _DSPARK_ADP_WIRE_TRAILER_LEN, + _DSPARK_EXACT_NATIVE_YIELD_INDEX, + _DSPARK_EXACT_WIRE_PREFIX_LEN, + PyExecutor, + _classify_dspark_exact_generation_rows, + _decode_dspark_exact_expected_yields, + _dspark_exact_common_graph_batch_size, + _dspark_exact_secondary_padding_ready, + _encode_dspark_exact_expected_yields, + _validate_dspark_adp_acceptance_gate, + _validate_dspark_exact_bucket, +) + + +class _Batch: + def __init__(self, generation_requests): + self.generation_requests = generation_requests + self.context_requests = [] + + @property + def batch_size(self): + return len(self.generation_requests) + + @property + def num_context_requests(self): + return 0 + + @property + def num_generation_requests(self): + return len(self.generation_requests) + + +def _request(*, adp_dummy=False, other_dummy=False): + return SimpleNamespace( + is_dummy=bool(adp_dummy or other_dummy), + is_attention_dp_dummy=bool(adp_dummy), + ) + + +def _payload(*, rows, native, compact): + payload = [0] * (_DSPARK_EXACT_WIRE_PREFIX_LEN + len(compact) + _DSPARK_ADP_WIRE_TRAILER_LEN) + payload[1] = rows + payload[_DSPARK_EXACT_NATIVE_YIELD_INDEX] = native + payload[_DSPARK_EXACT_WIRE_PREFIX_LEN : _DSPARK_EXACT_WIRE_PREFIX_LEN + len(compact)] = compact + return payload + + +def _agreement(batch, peer_sizes=(3, 4)): + return ADPShapeAgreement( + iteration=9, + batch_identity=id(batch), + local_batch_size=batch.batch_size, + peer_batch_sizes=peer_sizes, + all_can_graph=True, + widest_batch_size=max(peer_sizes), + graph_batch_size=4, + draft_len=5, + padding_ready=True, + ragged_bucket=16, + reuse_graph_shape=True, + ) + + +def test_exact_yield_encoding_is_conservative_and_pads_missing_cells(): + exact_local = SimpleNamespace( + native_expected_yield=10.0, + compact_expected_yields=(9.25, float("nan")), + ) + + native, compact = _encode_dspark_exact_expected_yields( + exact_local=exact_local, num_cells=3, yield_scale=100 + ) + + assert native == 1000 + assert compact == (925, -1, -1) + + +def test_exact_yield_encoding_rejects_an_invalid_scale(): + with pytest.raises(ValueError, match="scale must be positive"): + _encode_dspark_exact_expected_yields(exact_local=None, num_cells=1, yield_scale=0) + + +def test_exact_yield_decode_aggregates_yield_and_worst_rank_loss(): + cells = ((128, 512), (128, 640)) + payloads = [ + _payload( + rows=2, + native=10_000_000, + compact=(8_000_000, 9_000_000), + ), + _payload( + rows=4, + native=20_000_000, + compact=(18_000_000, 14_000_000), + ), + _payload(rows=0, native=0, compact=(0, 0)), + ] + + native, compact, max_loss = _decode_dspark_exact_expected_yields( + payloads=payloads, + exact_cells=cells, + graph_batch_size=128, + yield_scale=1_000_000, + ) + + assert native == 30.0 + assert compact == {512: 26.0, 640: 23.0} + assert max_loss == {512: 1.0, 640: 1.5} + + +def test_exact_yield_decode_fails_closed_on_an_infeasible_peer(): + payloads = [ + _payload(rows=2, native=10_000_000, compact=(8_000_000,)), + _payload(rows=4, native=20_000_000, compact=(-1,)), + ] + + _, compact, max_loss = _decode_dspark_exact_expected_yields( + payloads=payloads, + exact_cells=((128, 512),), + graph_batch_size=128, + yield_scale=1_000_000, + ) + + assert compact == {512: 0.0} + assert max_loss == {512: 0.0} + + +def test_exact_row_classifier_authenticates_only_the_idle_adp_dummy(): + real = [_request(), _request()] + assert _classify_dspark_exact_generation_rows(real) == (real, False) + assert _classify_dspark_exact_generation_rows([]) == (None, False) + assert _classify_dspark_exact_generation_rows([_request(adp_dummy=True)]) == ( + [], + True, + ) + assert _classify_dspark_exact_generation_rows([real[0], _request(adp_dummy=True)]) == ( + None, + False, + ) + assert _classify_dspark_exact_generation_rows([_request(other_dummy=True)]) == (None, False) + + +def test_exact_common_graph_size_preserves_all_rank_termination(): + def round_up(rows): + return 16 if rows <= 16 else 32 + + assert _dspark_exact_common_graph_batch_size([[0, 16], [0, 0]], round_up) == 16 + + def reject_round_up(_rows): + raise AssertionError("all-zero peers must not perform a graph lookup") + + assert _dspark_exact_common_graph_batch_size([[0, 0], [0, 0]], reject_round_up) == 0 + + +@pytest.mark.parametrize( + ("budget", "secondary_ready", "expected"), + [(48, False, True), (49, False, True), (50, False, False), (50, True, True)], +) +def test_secondary_dummy_is_required_only_for_a_multirow_remainder( + budget, secondary_ready, expected +): + payloads = [[0, 16, 0], [0, 0, int(secondary_ready)]] + assert ( + _dspark_exact_secondary_padding_ready( + payloads, + graph_batch_size=16, + verifier_budget=budget, + secondary_ready_index=2, + ) + is expected + ) + + +def test_exact_bucket_must_match_the_selected_graph_key(): + _validate_dspark_exact_bucket(exact_shape=(16, 48, 3), bucket=48) + _validate_dspark_exact_bucket(exact_shape=None, bucket=None) + with pytest.raises(RuntimeError, match="selected captured verifier bucket"): + _validate_dspark_exact_bucket(exact_shape=(16, 48, 3), bucket=47) + + +def test_attention_dp_rejects_a_rank_local_acceptance_gate(): + with pytest.raises(ValueError, match="acceptance_rate_window_size"): + _validate_dspark_adp_acceptance_gate( + confidence_enabled=True, + attention_dp_enabled=True, + speculation_gate_enabled=True, + ) + _validate_dspark_adp_acceptance_gate( + confidence_enabled=True, + attention_dp_enabled=False, + speculation_gate_enabled=True, + ) + + +@pytest.mark.parametrize( + ("peer_sizes", "expected"), + [((3, 4), (True, True)), ((3, 0), (False, True))], +) +def test_can_queue_reuses_policy_agreement_without_another_collective(peer_sizes, expected): + batch = _Batch([object()] * 3) + + def reject_allgather(_value): + raise AssertionError("cached agreement issued another collective") + + executor = PyExecutor.__new__(PyExecutor) + executor.enable_attention_dp = True + executor.dist = SimpleNamespace(tp_size=8, tp_allgather=reject_allgather) + executor.model_engine = SimpleNamespace( + _dspark_confidence_enabled=True, + cuda_graph_runner=SimpleNamespace(adp_shape_agreement=_agreement(batch, peer_sizes)), + ) + executor.kv_connector_manager = None + executor.kv_cache_transceiver = None + executor.drafter = None + executor.iter_counter = 9 + executor._dspark_dynamic_handled_signature = (9, id(batch), 3, 0, 3) + executor._handle_dynamic_draft_len = lambda _batch: None + + assert executor._can_queue(batch) == expected + + +def _draft_len_executor(*, confidence_enabled): + spec_mode = SimpleNamespace( + support_dynamic_draft_len=lambda: False, + is_pard=lambda: False, + ) + spec_config = SimpleNamespace( + draft_len_schedule=None, + spec_dec_mode=spec_mode, + enable_confidence_scheduling=confidence_enabled, + use_rejection_sampling=False, + is_linear_tree=True, + ) + executor = PyExecutor.__new__(PyExecutor) + executor.model_engine = SimpleNamespace( + max_draft_len=5, + max_total_draft_tokens=5, + runtime_draft_len=None, + spec_config=spec_config, + cuda_graph_runner=SimpleNamespace(adp_shape_agreement=None), + ) + executor.speculation_permanently_disabled = False + executor.iter_counter = 9 + return executor + + +def test_fixed_draft_len_keeps_static_path_without_dynamic_signature(): + executor = _draft_len_executor(confidence_enabled=False) + batch = _Batch([object()]) + + executor._handle_dynamic_draft_len(batch) + executor._handle_dynamic_draft_len(batch) + + assert executor.model_engine.runtime_draft_len == 5 + assert not hasattr(executor, "_dspark_dynamic_handled_signature") + + +def test_confidence_draft_len_runs_only_once_per_iteration(): + executor = _draft_len_executor(confidence_enabled=True) + request = SimpleNamespace( + py_draft_tokens=[1, 2, 3, 4, 5], + py_needs_onehot_draft_probs=False, + ) + batch = _Batch([request]) + calls = 0 + + def select_draft_len(_batch): + nonlocal calls + calls += 1 + return 3 + + executor._dspark_confidence_draft_len = select_draft_len + executor._handle_dynamic_draft_len(batch) + executor._handle_dynamic_draft_len(batch) + + assert calls == 1 + assert executor.model_engine.runtime_draft_len == 3 + assert request.py_draft_tokens == [1, 2, 3] + assert executor._dspark_dynamic_handled_signature == (9, id(batch), 1, 0, 1) + + +def test_rebalance_suspends_both_dummy_variants_and_releases_the_pair_once(): + low = SimpleNamespace(py_request_id=10) + high = SimpleNamespace(py_request_id=11) + suspended_ids = [] + resumed_ids = [] + + class _Manager: + def is_request_active(self, request_id): + return request_id in {10, 11} + + def suspend_request(self, request): + suspended_ids.append(request.py_request_id) + + def resume_request(self, request): + resumed_ids.append(request.py_request_id) + return request.py_request_id != 11 + + released = [] + runner = SimpleNamespace( + padding_dummy_requests={5: low}, + secondary_padding_dummy_requests={5: high}, + release_padding_dummy=lambda resource, draft_len: released.append((resource, draft_len)), + ) + executor = PyExecutor.__new__(PyExecutor) + executor.model_engine = SimpleNamespace(cuda_graph_runner=runner) + executor.resource_manager = object() + manager = _Manager() + + suspended = executor._suspend_padding_dummies_for_rebalance(manager) + executor._resume_padding_dummies_after_rebalance(manager, suspended) + + assert suspended == [(5, 0, low), (5, 1, high)] + assert suspended_ids == [10, 11] + assert resumed_ids == [10, 11] + assert released == [(executor.resource_manager, 5)] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_fused_attention_route.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_fused_attention_route.py new file mode 100644 index 000000000000..7b1a6e030a4b --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_fused_attention_route.py @@ -0,0 +1,328 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU-only DSpark routing contracts; no Torch, CuTe, or CUDA runtime is loaded. + +Execute the production Python guards with tensor metadata and mocked launchers. +The AST extraction follows the existing K2/K3 routing-test prototype. Scalar +mask checks are not a substitute for compiling and testing the GPU kernels. +""" + +import ast +import math +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +_ROOT = Path(__file__).resolve().parents[5] +_OPS = _ROOT / "tensorrt_llm/_torch/custom_ops" +_ATTENTION = _OPS / "dspark_attention_custom_op.py" +_PREPARATION = _OPS / "dspark_rmsnorm_rope_custom_op.py" +_KERNELS = _ROOT / "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark" +_MODEL = _ROOT / "tensorrt_llm/_torch/models/modeling_dspark.py" +_PHYSICAL_KS = tuple(range(1, 9)) + + +def _tree(path): + return ast.parse(path.read_text(encoding="utf-8")) + + +def _load(path, name, namespace, owner=None): + nodes = _tree(path).body + if owner is not None: + nodes = next(node for node in nodes if getattr(node, "name", None) == owner).body + node = next(node for node in nodes if getattr(node, "name", None) == name) + node.decorator_list = [] + future = ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0) + module = ast.fix_missing_locations(ast.Module(body=[future, node], type_ignores=[])) + exec(compile(module, str(path), "exec"), namespace) + return namespace[name] + + +def _constants(path): + return { + node.targets[0].id: ast.literal_eval(node.value) + for node in _tree(path).body + if isinstance(node, ast.Assign) + and isinstance(node.targets[0], ast.Name) + and node.targets[0].id.startswith("_DSV4_DSPARK_") + } + + +def _tensor(shape, dtype="bf16", *, cuda=True, contiguous=True, strides=None): + if strides is None: + strides = tuple(math.prod(shape[i + 1 :]) for i in range(len(shape))) + return SimpleNamespace( + shape=shape, + ndim=len(shape), + dtype=dtype, + is_cuda=cuda, + is_contiguous=lambda: contiguous, + numel=lambda: math.prod(shape), + stride=lambda dim=None: strides if dim is None else strides[dim], + data_ptr=lambda: 0, + ) + + +def _guards(sm=100): + namespace = { + "torch": SimpleNamespace(bfloat16="bf16", float32="fp32", int32="i32", int64="i64"), + "get_sm_version": lambda: sm, + "_log_unsupported": lambda *args: False, + } + namespace.update(_constants(_PREPARATION)) + for name in ( + "_get_dspark_arch_str", + "is_fused_dspark_rmsnorm_rope_supported", + "_is_dspark_cache_write_supported", + "_is_dspark_draft_block_supported", + "is_fused_dspark_attention_preparation_supported", + ): + _load(_PREPARATION, name, namespace) + namespace.update(_constants(_ATTENTION)) + for name in ( + "is_dsv4_dspark_attention_config_supported", + "is_fused_dsv4_dspark_attention_supported", + "_is_fused_dsv4_dspark_attention_input_supported", + ): + _load(_ATTENTION, name, namespace) + return namespace + + +def _inputs(k, batch=2): + return [ + _tensor((batch, k, 128, 512)), + _tensor((batch, 8, 512)), + _tensor((batch + 1, 128, 512), contiguous=False, strides=(196608, 512, 1)), + _tensor((batch,), "i32"), + _tensor((batch,), "i32"), + _tensor((batch,), "i64"), + _tensor((128,), "fp32"), + _tensor((batch, k, 32, 2), "fp32"), + ] + + +class TestDSparkFusedAttentionRoute(unittest.TestCase): + def test_storage_bounded_range_is_consistent(self): + for path in (_ATTENTION, _PREPARATION): + constants = _constants(path) + self.assertEqual(constants["_DSV4_DSPARK_BLOCK_SIZES"], _PHYSICAL_KS) + self.assertEqual(constants["_DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE"], 8) + kernel = next( + node + for node in _tree(_KERNELS / "attention.py").body + if isinstance(node, ast.ClassDef) and node.name == "DSparkAttention" + ) + page_size = next( + ast.literal_eval(node.value) + for node in kernel.body + if isinstance(node, ast.Assign) and node.targets[0].id == "page_size_draft" + ) + self.assertEqual(_PHYSICAL_KS, tuple(range(1, page_size + 1))) + + def test_config_routes_every_k_and_rejects_out_of_range(self): + supported = _guards()["is_dsv4_dspark_attention_config_supported"] + for k in (-1, 0, *_PHYSICAL_KS, 9, 16, 128): + with self.subTest(k=k): + self.assertEqual(supported(k, 128, 512, 128), k in _PHYSICAL_KS) + for k in _PHYSICAL_KS: + self.assertFalse(supported(k, 64, 512, 128)) + self.assertFalse(supported(k, 128, 256, 128)) + self.assertFalse(supported(k, 128, 512, 64)) + + def test_kernel_constructor_keeps_geometry_and_storage_guards(self): + class KernelBase: + def __init__(self, *args, **kwargs): + self.seq_len_q = kwargs["seq_len_q"] + + namespace = { + "cutlass": SimpleNamespace(Float32="fp32"), + "cute": SimpleNamespace(jit=lambda f: f), + "DSparkAttentionKernel": KernelBase, + } + kernel = _load(_KERNELS / "attention.py", "DSparkAttention", namespace) + args = ("fp32", (128, 128), (128, 256), 1, 8, 128, 0.0) + for k in _PHYSICAL_KS: + with self.subTest(k=k): + op = kernel(*args, arch_str="sm_100", seq_len_q=k, inverse_rope_dim=64) + self.assertEqual(op.seq_len_q, k) + self.assertEqual(op.fixed_cache_seq_len, 128 + k) + self.assertEqual(op.tma_page_size_draft, 128) + for k in (-1, 0, 9, 16): + with self.subTest(k=k), self.assertRaises(ValueError): + kernel(*args, arch_str="sm_100", seq_len_q=k) + with self.assertRaises(ValueError): + kernel(*args, arch_str="sm_100", seq_len_q=8, inverse_rope_dim=32) + with self.assertRaises(ValueError): + kernel(*args[:4], 16, *args[5:], arch_str="sm_100", seq_len_q=8) + + def test_attention_tensor_contract_for_every_k_and_architecture(self): + for sm in (90, 100, 101, 103, 109, 120, 121): + gate = _guards(sm)["_is_fused_dsv4_dspark_attention_input_supported"] + for k in (0, *_PHYSICAL_KS, 9, 16): + with self.subTest(sm=sm, k=k): + self.assertEqual(gate(*_inputs(k)), sm in (100, 103) and k in _PHYSICAL_KS) + + def test_attention_dtype_shape_and_layout_guards_survive(self): + gate = _guards()["_is_fused_dsv4_dspark_attention_input_supported"] + bad_inputs = ( + (0, _tensor((2, 4, 128, 512), dtype="fp32")), + (0, _tensor((2, 4, 64, 512))), + (0, _tensor((2, 4, 128, 256))), + (0, _tensor((2, 4, 512))), + (1, _tensor((2, 7, 512))), + (1, _tensor((2, 9, 512))), + (1, _tensor((2, 8, 512), dtype="fp32")), + (2, _tensor((3, 64, 512))), + (2, _tensor((3, 128, 512), dtype="fp32")), + (2, _tensor((3, 128, 512), strides=(0, 512, 1))), + (2, _tensor((3, 128, 512), strides=(196609, 512, 1))), + (2, _tensor((3, 128, 512), strides=(196608, 513, 1))), + (2, _tensor((3, 128, 512), strides=(196608, 512, 2))), + (3, _tensor((2,), "i64")), + (4, _tensor((2,), "i64")), + (4, _tensor((3,), "i32")), + (5, _tensor((2,), "i32")), + (5, _tensor((3,), "i64")), + (6, _tensor((128,), "bf16")), + (6, _tensor((64,), "fp32")), + (7, _tensor((2, 4, 31, 2), "fp32")), + ) + for index, bad in bad_inputs: + with self.subTest(index=index, shape=bad.shape, dtype=bad.dtype): + inputs = _inputs(4) + inputs[index] = bad + self.assertFalse(gate(*inputs)) + for k in _PHYSICAL_KS: + for index in range(8): + inputs = _inputs(k) + inputs[index].is_cuda = False + self.assertFalse(gate(*inputs)) + if index != 2: # The worker's strided cache-page view is supported. + inputs = _inputs(k) + inputs[index].is_contiguous = lambda: False + self.assertFalse(gate(*inputs)) + + def test_preparation_bounds_and_guards_for_every_k(self): + gate = _guards()["is_fused_dspark_attention_preparation_supported"] + for k in (0, *_PHYSICAL_KS, 9, 16): + with self.subTest(k=k): + inputs = [ + _tensor((2, 1, 512)), + _tensor((2, k, 512)), + _tensor((512,)), + _tensor((2, 32, 2), "fp32"), + _tensor((2 * k, 32, 2), "fp32"), + _inputs(k)[2], + _tensor((2,), "i64"), + _tensor((2,), "i64"), + ] + self.assertEqual(gate(*inputs), k in _PHYSICAL_KS) + if k in _PHYSICAL_KS: + for index in range(len(inputs)): + original = inputs[index].is_cuda + inputs[index].is_cuda = False + self.assertFalse(gate(*inputs)) + inputs[index].is_cuda = original + inputs[4] = _tensor((2 * k + 1, 32, 2), "fp32") + self.assertFalse(gate(*inputs)) + + def test_masks_cover_short_full_wrapped_and_empty_windows(self): + namespace = { + "cutlass": SimpleNamespace(Int32=int, const_expr=lambda x: x), + "cute": SimpleNamespace(elem_less=lambda a, b: a < b), + } + path = _KERNELS / "attention_kernel.py" + mask = _load(path, "is_score_valid", namespace, "DSparkAttentionKernel") + valid_length = _load(path, "get_window_valid_len", namespace, "DSparkAttentionKernel") + kernel = SimpleNamespace( + mma_qk_tiler=(128, 128, 128), + window_valid_len_from_tensor=True, + tma_page_size_win=128, + ) + for k in _PHYSICAL_KS: + for end in (0, 5, 127, 128, 257, 390): + for valid in (-1, 0, 1, 3, 128, 129): + with self.subTest(k=k, end=end, valid=valid): + length = valid_length(kernel, [valid], 0) + expected = {(end - age) % 128 for age in range(min(max(valid, 0), 128))} + actual = { + col for col in range(128) if mask(kernel, col, 0, 128 + k, length, end) + } + self.assertEqual(actual, expected) + draft = { + col for col in range(128) if mask(kernel, col, 1, 128 + k, length, end) + } + self.assertEqual(draft, set(range(k))) + + def test_cache_keys_separate_k_and_arch_but_not_runtime_shapes(self): + namespace = _guards() + cache = {} + compile_kernel = Mock(side_effect=lambda *args: Mock()) + namespace["_dspark_attention_kernel_cache"] = cache + namespace["_compile_dspark_attention"] = compile_kernel + capturing = Mock(return_value=False) + namespace["torch"].cuda = SimpleNamespace(is_current_stream_capturing=capturing) + namespace["torch"].empty_like = lambda q: _tensor(q.shape) + run = _load(_ATTENTION, "_run_dspark_attention", namespace) + for arch in ("sm_100", "sm_103"): + namespace["_get_dspark_arch_str"] = lambda arch=arch: arch + for k in _PHYSICAL_KS: + with self.subTest(k=k, arch=arch): + for batch, scale in ((1, 0.125), (3, 0.25)): + self.assertEqual(run(*_inputs(k, batch), scale).shape, (batch, k, 128, 512)) + self.assertEqual(cache[(k, arch)].call_count, 2) + self.assertEqual(compile_kernel.call_count, 16) + self.assertEqual(capturing.call_count, 16) + self.assertEqual(len({id(compiled) for compiled in cache.values()}), 16) + + def test_graph_capture_only_uses_warmed_k(self): + for k in _PHYSICAL_KS: + namespace = _guards() + namespace["_dspark_attention_kernel_cache"] = {} + namespace["_compile_dspark_attention"] = Mock(side_effect=AssertionError("compile")) + namespace["torch"].cuda = SimpleNamespace(is_current_stream_capturing=lambda: True) + namespace["torch"].empty_like = Mock(side_effect=AssertionError("allocation")) + run = _load(_ATTENTION, "_run_dspark_attention", namespace) + with self.subTest(k=k), self.assertRaisesRegex(RuntimeError, "warmed up"): + run(*_inputs(k), 0.125) + + def test_unsupported_k_warmup_does_not_allocate(self): + namespace = _guards() + warmup = _load(_ATTENTION, "warmup_fused_dsv4_dspark_attention", namespace) + # The metadata-only torch object cannot allocate; early return is required. + for k in (-1, 0, 9, 16): + warmup(k, 1e-6) + + def test_query_storage_and_fallback_routes_remain_separate(self): + attention = ast.unparse(_tree(_KERNELS / "attention.py")) + self.assertIn("(self.num_heads, self.head_dim, self.seq_len_q, batch_size)", attention) + self.assertIn("batch_size * self.seq_len_q * self.inverse_rope_dim", attention) + self.assertIn("self.fixed_cache_seq_len = self.window_size + seq_len_q", attention) + model = {} + _load(_MODEL, "dspark_attention_forward_batched", model) + function = next( + node + for node in _tree(_MODEL).body + if getattr(node, "name", None) == "dspark_attention_forward_batched" + ) + route = next( + node + for node in function.body + if isinstance(node, ast.If) + and ast.unparse(node.test) == "use_fused_dsv4_dspark_attention" + ) + fused = ast.unparse(ast.Module(body=route.body, type_ignores=[])) + fallback = ast.unparse(ast.Module(body=route.orelse, type_ignores=[])) + self.assertIn("cute_dsl_dspark_rmsnorm_rope_draft_block(", fused) + self.assertIn("fused_dsv4_dspark_attention(", fused) + self.assertNotIn("fused_dsv4_dspark_attention(", fallback) + self.assertIn( + "get_dspark_topk_idxs_batched(window_size, block, start_pos, valid_len)", fallback + ) + self.assertIn("torch.cat([cache_rows, kv], dim=1)", fallback) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_evidence.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_evidence.py new file mode 100644 index 000000000000..3fca02f00a02 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_evidence.py @@ -0,0 +1,458 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only, deliberately synthetic protocol fixtures; no production evidence.""" + +import copy +import dataclasses +import hashlib +import json +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace + +from tensorrt_llm._torch.speculative import dspark_mixed_evidence as mixed +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig, DSparkMixedEvidenceConfig + + +def sha(raw): + return hashlib.sha256(raw).hexdigest() + + +def encode(value): + return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() + + +def identity(): + return mixed.MixedRuntimeIdentity(*("abcdef"[index] * 64 for index in range(6))) + + +def geometry(): + # Rank 0 prefills 64 rows; rank 1 decodes with native 32 versus compact 16. + # Every bootstrap/dummy generation is still charged its full K+1 rows. + features = ( + (64, 4096, 8192, 0, 0, 0, 0, 1, 0, 1, 0, 0), + (0, 0, 0, 0, 1024, 128, 0, 0, 8, 0, 8, 0), + ) + ((0,) * 12,) * 6 + return mixed.MixedGeometry(features, 96, (64, 16, 0, 0, 0, 0, 0, 0), (64, 32, 0, 0, 0, 0, 0, 0)) + + +def fixture(): + live = identity() + table = { + "schema": "dspark-mixed-breakable-costs-v1", + "identity": live.payload(), + "protocol_sha256": "1" * 64, + "action_semantics": mixed._ACTION, + "pricing_semantics": mixed._PRICE_SEMANTICS, + "synthetic": False, + "provenance": "closed_measured_breakable_mixed", + "minimum_predicted_gain": 0.01, + "cells": [ + { + "geometry": geometry().payload(), + "native_lower_ms": 7.0, + "candidate_upper_ms": 6.9, + "measurement_result_sha256": "2" * 64, + } + ], + } + receipt = { + "schema": "dspark-mixed-breakable-admission-v1", + "identity": live.payload(), + "costs_sha256": sha(encode(table)), + "protocol_sha256": "1" * 64, + "action_semantics": mixed._ACTION, + "pricing_semantics": mixed._PRICE_SEMANTICS, + "production_admission": True, + "source_and_measurements_verified": True, + "accuracy_pass": True, + "runtime_invariants_pass": True, + "matched_workload_pass": True, + "accuracy_result_sha256": "3" * 64, + "runtime_invariants_result_sha256": "4" * 64, + "matched_workload_result_sha256": "5" * 64, + "selector_replay_sha256": "6" * 64, + "measurement_results_sha256": sha(encode(["2" * 64])), + "economics": { + "policy_steps": 10, + "compact_choices": 4, + "gross_compact_value_ms_lower_bound": 3.0, + "fixed_confidence_path_overhead_ms_upper_bound": 0.1, + "safety_margin_ms": 0.5, + }, + } + return table, receipt, live + + +def resolve(table, receipt, live=None, *, update_receipt=True): + table_raw = encode(table) + receipt = copy.deepcopy(receipt) + if update_receipt: + receipt["costs_sha256"] = sha(table_raw) + receipt_raw = encode(receipt) + binding = mixed.MixedEvidenceBinding(sha(table_raw), sha(receipt_raw), "1" * 64) + return mixed.resolve_mixed_evidence( + binding, costs_raw=table_raw, admission_raw=receipt_raw, live_identity=live or identity() + ) + + +class TestMixedEvidence(unittest.TestCase): + def assert_native(self, table, receipt, live=None, **kwargs): + result = resolve(table, receipt, live, **kwargs) + self.assertIsNone(result.costs) + self.assertTrue(result.reason.startswith("mixed_evidence_invalid_native:")) + + def test_01_actual_binding_positive_synthetic_fixture_only(self): + table, receipt, live = fixture() + resolved = resolve(table, receipt, live) + self.assertEqual(resolved.reason, "mixed_evidence_bound") + self.assertEqual(resolved.costs.lookup(geometry()).candidate_upper_ms, 6.9) + self.assertEqual(resolved.costs.minimum_predicted_gain, 0.01) + + def test_02_bad_fixed_restored_old_validation_schema(self): + table, receipt, live = fixture() + table["schema"] = "dspark-host-ragged-bcg-prices-v1" + self.assert_native(table, receipt, live) + table["schema"] = "dspark-mixed-breakable-costs-v1" + self.assertIsNotNone(resolve(table, receipt, live).costs) + table["schema"] = "dspark-host-ragged-bcg-prices-v1" + self.assert_native(table, receipt, live) + + def test_03_absent_optional_binding_does_not_read_payloads(self): + result = mixed.resolve_mixed_evidence( + None, costs_raw=None, admission_raw=None, live_identity=identity() + ) + self.assertEqual(result.reason, "mixed_evidence_absent_native") + self.assertIsNone(result.costs) + + def test_04_decode_policy_object_and_flag_unchanged_for_all_outcomes(self): + decode = {"enabled": True, "identity": "existing-decode", "table": object()} + before = dict(decode) + table, receipt, live = fixture() + for binding in (None, mixed.MixedEvidenceBinding("0" * 64, "1" * 64, "2" * 64)): + mixed.resolve_mixed_evidence( + binding, costs_raw=b"bad", admission_raw=b"bad", live_identity=live + ) + self.assertEqual(decode, before) + self.assertTrue(decode["enabled"]) + self.assertIsNotNone(resolve(table, receipt, live).costs) + self.assertEqual(decode, before) + + def test_05_live_identity_is_not_read_from_table(self): + table, receipt, live = fixture() + for field in ( + "source_tree_sha256", + "runtime_snapshot_sha256", + "model_config_sha256", + "gpu_class_sha256", + "topology_sha256", + "selector_identity_sha256", + ): + with self.subTest(field=field): + self.assert_native(table, receipt, dataclasses.replace(live, **{field: "9" * 64})) + + def test_06_invalid_supported_runtime_and_mutable_grid_rejected(self): + live = identity() + for overrides in ( + {"physical_k": 2}, + {"physical_k": True}, + {"world_size": 1}, + {"execution_mode": "FULL"}, + {"body_buckets": list(live.body_buckets)}, + {"body_buckets": (96, 128, 192, 256)}, + ): + with self.subTest(overrides=overrides), self.assertRaises(ValueError): + dataclasses.replace(live, **overrides) + + def test_07_payloads_are_deeply_frozen_into_typed_tuples(self): + table, receipt, live = fixture() + costs = resolve(table, receipt, live).costs + table["cells"][0]["geometry"]["rank_features"][1]["eligible_generations"] = 0 + table["cells"][0]["native_lower_ms"] = 0 + self.assertEqual(costs.lookup(geometry()).native_lower_ms, 7.0) + with self.assertRaises(dataclasses.FrozenInstanceError): + costs.cells[0].candidate_upper_ms = 1 + + def test_08_uncovered_geometry_is_native_no_interpolation(self): + table, receipt, live = fixture() + costs = resolve(table, receipt, live).costs + unseen = dataclasses.replace(geometry(), actual_rank_rows=(64, 15, 0, 0, 0, 0, 0, 0)) + self.assertIsNone(costs.lookup(unseen)) + self.assertIsNone(costs.lookup(geometry().payload())) + + def test_09_duplicate_geometry_rejects_favorable_reselection(self): + table, receipt, _ = fixture() + table["cells"].append(copy.deepcopy(table["cells"][0])) + table["cells"][1]["candidate_upper_ms"] = 1 + self.assert_native(table, receipt) + + def test_10_cell_list_empty_or_unbounded_rejected(self): + table, receipt, _ = fixture() + for cells in ([], [table["cells"][0]] * 257): + bad = copy.deepcopy(table) + bad["cells"] = cells + self.assert_native(bad, receipt) + + def test_11_context_and_fixed_generation_rows_cannot_be_discounted(self): + original = geometry() + for change in ( + {"actual_rank_rows": (63, 16, 0, 0, 0, 0, 0, 0)}, + {"actual_rank_rows": (64, 7, 0, 0, 0, 0, 0, 0)}, + {"native_rank_rows": (64, 31, 0, 0, 0, 0, 0, 0)}, + ): + with self.subTest(change=change), self.assertRaises(ValueError): + dataclasses.replace(original, **change) + features = list(original.rank_features) + features[1] = (0, 0, 0, 0, 1024, 128, 1, 0, 7, 0, 8, 0) + with self.assertRaises(ValueError): + dataclasses.replace(original, rank_features=tuple(features)) + + def test_12_decode_only_and_noncompact_geometry_rejected(self): + original = geometry() + features = list(original.rank_features) + features[0] = (64, 4096, 8192, 0, 0, 0, 0, 0, 0, 1, 0, 0) + with self.assertRaises(ValueError): + dataclasses.replace(original, rank_features=tuple(features)) + with self.assertRaises(ValueError): + dataclasses.replace(original, actual_rank_rows=original.native_rank_rows) + + def test_13_rank_order_is_not_coalesced(self): + table, receipt, live = fixture() + costs = resolve(table, receipt, live).costs + original = geometry() + reordered = dataclasses.replace( + original, + rank_features=original.rank_features[::-1], + actual_rank_rows=original.actual_rank_rows[::-1], + native_rank_rows=original.native_rank_rows[::-1], + ) + self.assertIsNone(costs.lookup(reordered)) + + def test_14_unknown_null_and_validation_only_fields_reject(self): + table, receipt, _ = fixture() + for field, value in ( + ("validation_only", False), + ("pricing_sequence", [96]), + ("confidence_sps_table_path", "old.json"), + ("unknown", None), + ): + bad = copy.deepcopy(table) + bad[field] = value + self.assert_native(bad, receipt) + + def test_15_synthetic_prices_and_claim_only_receipts_reject(self): + table, receipt, _ = fixture() + for field in ( + "production_admission", + "source_and_measurements_verified", + "accuracy_pass", + "runtime_invariants_pass", + "matched_workload_pass", + ): + bad = copy.deepcopy(receipt) + bad[field] = False + self.assert_native(table, bad) + table["synthetic"] = True + self.assert_native(table, receipt) + + def test_16_proxy_is_never_relabelled_as_whole_step(self): + table, receipt, _ = fixture() + for key in ("table", "receipt"): + bad_table, bad_receipt = copy.deepcopy(table), copy.deepcopy(receipt) + (bad_table if key == "table" else bad_receipt)["pricing_semantics"] = ( + "qualified_group_whole_step" + ) + self.assert_native(bad_table, bad_receipt) + + def test_17_protocol_action_and_measurement_binding_reject_drift(self): + table, receipt, _ = fixture() + for key, value in ( + ("protocol_sha256", "9" * 64), + ("action_semantics", "uniform-L"), + ("measurement_results_sha256", "9" * 64), + ): + bad = copy.deepcopy(receipt) + bad[key] = value + self.assert_native(table, bad) + + def test_18_all_result_hashes_required(self): + table, receipt, _ = fixture() + for field in ( + "accuracy_result_sha256", + "runtime_invariants_result_sha256", + "matched_workload_result_sha256", + "selector_replay_sha256", + "measurement_results_sha256", + ): + bad = copy.deepcopy(receipt) + bad[field] = "not-a-hash" + self.assert_native(table, bad) + + def test_19_positive_economics_recomputed_no_boolean_override(self): + table, receipt, _ = fixture() + for change in ( + {"compact_choices": 0}, + {"gross_compact_value_ms_lower_bound": 1.0}, + {"policy_steps": True}, + {"safety_margin_ms": -1.0}, + {"compact_choices": 11}, + ): + bad = copy.deepcopy(receipt) + bad["economics"].update(change) + self.assert_native(table, bad) + + def test_20_no_relaxation_of_one_percent_immediate_gate(self): + table, receipt, _ = fixture() + for margin in (0.0, 0.009, True): + table["minimum_predicted_gain"] = margin + self.assert_native(table, receipt) + table["minimum_predicted_gain"] = 0.02 + self.assertEqual(resolve(table, receipt).costs.minimum_predicted_gain, 0.02) + + def test_21_nonpositive_and_nonfinite_costs_rejected(self): + table, receipt, _ = fixture() + for field in ("native_lower_ms", "candidate_upper_ms"): + for value in (0, -1, True, "7.0"): + bad = copy.deepcopy(table) + bad["cells"][0][field] = value + self.assert_native(bad, receipt) + raw = encode(table).replace(b'"candidate_upper_ms":6.9', b'"candidate_upper_ms":1e9999') + receipt["costs_sha256"] = sha(raw) + self.assert_raw_native(raw, encode(receipt)) + + def assert_raw_native(self, table_raw, receipt_raw): + binding = mixed.MixedEvidenceBinding(sha(table_raw), sha(receipt_raw), "1" * 64) + result = mixed.resolve_mixed_evidence( + binding, costs_raw=table_raw, admission_raw=receipt_raw, live_identity=identity() + ) + self.assertIsNone(result.costs) + + def test_22_duplicate_nonfinite_malformed_deep_json_rejected(self): + _, receipt, _ = fixture() + for raw in ( + b'{"schema":"x","schema":"y"}', + b'{"x":NaN}', + b'{"x":Infinity}', + b"{malformed", + b"\xff", + b"[" * 1500 + b"]" * 1500, + ): + self.assert_raw_native(raw, encode(receipt)) + + def test_23_missing_raw_oversized_and_digest_drift_rejected(self): + table, receipt, live = fixture() + raw, qraw = encode(table), encode(receipt) + binding = mixed.MixedEvidenceBinding(sha(raw), sha(qraw), "1" * 64) + for value in (None, raw + b" ", bytearray(raw), b"", b" " * (mixed._MAX_BYTES + 1)): + result = mixed.resolve_mixed_evidence( + binding, costs_raw=value, admission_raw=qraw, live_identity=live + ) + self.assertIsNone(result.costs) + + def test_24_table_receipt_cross_binding_rejected(self): + table, receipt, _ = fixture() + table["cells"][0]["candidate_upper_ms"] = 6.8 + self.assert_native(table, receipt, update_receipt=False) + + def test_25_rank_extents_strict_types_and_wire_limits(self): + original = geometry() + for value in (True, 1.0, -1, 2**63): + features = list(original.rank_features) + changed = list(features[0]) + changed[2] = value + features[0] = tuple(changed) + with self.subTest(value=value), self.assertRaises((TypeError, ValueError)): + dataclasses.replace(original, rank_features=tuple(features)) + + def test_26_decode_schema_is_not_mixed_evidence(self): + _, receipt, _ = fixture() + self.assert_raw_native(encode({"schema_version": 2}), encode(receipt)) + + def test_27_valid_bootstrap_reserves_full_k_plus_one_rows(self): + original = geometry() + features = list(original.rank_features) + features[1] = (0, 0, 0, 0, 1024, 128, 1, 0, 7, 0, 8, 4) + valid = dataclasses.replace( + original, rank_features=tuple(features), actual_rank_rows=(64, 20, 0, 0, 0, 0, 0, 0) + ) + self.assertEqual(valid.rank_features[1][-1], 4) + with self.assertRaises(ValueError): + dataclasses.replace(valid, actual_rank_rows=(64, 10, 0, 0, 0, 0, 0, 0)) + + def test_28_bad_fixed_restored_cost_pin_is_not_relabelled(self): + table, receipt, live = fixture() + raw, qraw = encode(table), encode(receipt) + bad_binding = mixed.MixedEvidenceBinding("0" * 64, sha(qraw), "1" * 64) + for binding, expected in ( + (bad_binding, False), + (mixed.MixedEvidenceBinding(sha(raw), sha(qraw), "1" * 64), True), + (bad_binding, False), + ): + result = mixed.resolve_mixed_evidence( + binding, costs_raw=raw, admission_raw=qraw, live_identity=live + ) + self.assertEqual(result.costs is not None, expected) + + def test_29_file_binding_requires_independent_live_fingerprint(self): + table, receipt, live = fixture() + table_raw, receipt_raw = encode(table), encode(receipt) + live_raw = encode( + {"schema": "dspark-mixed-live-fingerprint-v1", "identity": live.payload()} + ) + with tempfile.TemporaryDirectory() as directory: + paths = [Path(directory) / name for name in ("costs.json", "receipt.json", "live.json")] + for path, raw in zip(paths, (table_raw, receipt_raw, live_raw)): + path.write_bytes(raw) + config = SimpleNamespace( + costs_path=str(paths[0]), + costs_sha256=sha(table_raw), + admission_path=str(paths[1]), + admission_sha256=sha(receipt_raw), + protocol_sha256="1" * 64, + live_fingerprint_path=str(paths[2]), + live_fingerprint_sha256=sha(live_raw), + ) + self.assertIsNotNone(mixed.load_mixed_evidence(config).costs) + config.live_fingerprint_sha256 = "0" * 64 + self.assertIsNone(mixed.load_mixed_evidence(config).costs) + self.assertEqual(mixed.load_mixed_evidence(None).reason, "mixed_evidence_absent_native") + + def test_30_typed_config_rejects_invalid_pins_and_unknown_fields(self): + data = dict( + costs_path="costs.json", + costs_sha256="1" * 64, + admission_path="receipt.json", + admission_sha256="2" * 64, + protocol_sha256="3" * 64, + live_fingerprint_path="live.json", + live_fingerprint_sha256="4" * 64, + ) + self.assertEqual(DSparkMixedEvidenceConfig(**data).costs_path, "costs.json") + with self.assertRaises(ValueError): + DSparkMixedEvidenceConfig(**{**data, "costs_sha256": "invalid"}) + with self.assertRaises(ValueError): + DSparkMixedEvidenceConfig(**data, force=True) + + def test_31_feature_off_rejects_mixed_input_and_k1_clears_it(self): + config = SimpleNamespace( + enable_confidence_scheduling=False, + enable_fused_confidence_scheduler=False, + confidence_sts_path=None, + confidence_sps_table_path=None, + confidence_sps_live_fingerprint_path=None, + confidence_admission_receipt_path=None, + confidence_admission_receipt_sha256=None, + confidence_mixed_evidence=object(), + max_draft_len=3, + ) + with self.assertRaisesRegex(ValueError, "require enable_confidence_scheduling"): + DSparkDecodingConfig.validate_confidence_scheduling(config) + config.enable_confidence_scheduling, config.max_draft_len = True, 1 + self.assertIs(DSparkDecodingConfig.validate_confidence_scheduling(config), config) + self.assertFalse(config.enable_confidence_scheduling) + self.assertIsNone(config.confidence_mixed_evidence) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_runtime.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_runtime.py new file mode 100644 index 000000000000..73f8cb0aa618 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_mixed_runtime.py @@ -0,0 +1,350 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU contracts for optional DSpark mixed allocation; prices are synthetic.""" + +import copy +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor import dspark_mixed_runtime as runtime +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType +from tensorrt_llm._torch.speculative.dspark_mixed_evidence import ( + AdmittedMixedCosts, + MixedCostCell, + MixedEvidenceBinding, + MixedGeometry, + MixedRuntimeIdentity, +) +from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig +from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner + + +class _Event: + def __init__(self, ready=True): + self.ready = ready + self.records = 0 + + def query(self): + return self.ready + + def record(self): + self.records += 1 + + +def _request(request_id, *, context=False): + request = SimpleNamespace( + py_request_id=request_id, + py_seq_slot=request_id, + is_dummy=False, + py_batch_idx=0, + py_draft_tokens=[1, 2, 3], + py_verify_len=None, + state=SimpleNamespace(name="GENERATION_IN_PROGRESS"), + py_orig_prompt_len=128, + py_max_new_tokens=768, + context_chunk_size=64 if context else 0, + context_current_position=128 if context else 0, + py_num_compressed_tokens=0, + is_first_context_chunk=False, + ) + request.get_num_tokens = lambda beam: 128 + return request + + +def _group(monkeypatch, *, all_decode=False, stale=False): + monkeypatch.setattr(runtime, "_greedy", lambda _batch: True) + identity = MixedRuntimeIdentity(*(c * 64 for c in "abcdef")) + binding = MixedEvidenceBinding("1" * 64, "2" * 64, "3" * 64) + costs = AdmittedMixedCosts(identity, binding, 0.01, ()) + group = [] + for rank in range(8): + config = SimpleNamespace( + max_draft_len=3, + max_seq_len=4096, + max_num_tokens=4096, + enable_attention_dp=True, + torch_compile_enabled=False, + is_multimodal=False, + spec_config=SimpleNamespace( + max_draft_len=3, + block_size=3, + decoding_type="DSpark", + draft_is_embedded_in_target=True, + use_rejection_sampling=False, + ), + prefill_cuda_graph_backend=SimpleNamespace(name="BREAKABLE"), + prefill_cuda_graph_num_tokens=runtime.BODY_BUCKETS, + ) + generations = [_request(rank * 1000 + i + 1) for i in range(64)] + contexts = [_request(99999, context=True)] if rank == 0 and not all_decode else [] + batch = SimpleNamespace( + context_requests=contexts, + generation_requests=generations, + context_requests_last_chunk=contexts, + ) + planner = DSparkVerifyPlanner(cfg=DSparkScheduleConfig(3), device_windows=True) + planner._host_buffer = torch.full((64, 3), 5.0) + planner._host_confidence_stamp = torch.full((64,), 4, dtype=torch.int32) + planner._snapshot_valid, planner._copy_event = True, _Event() + slots = {r.py_request_id: i for i, r in enumerate(generations)} + incarnations = {r.py_request_id: i + 1 for i, r in enumerate(generations)} + worker = SimpleNamespace( + _draft_seq_host=5, + _req_to_slot=slots, + confidence_row_for=slots.__getitem__, + confidence_incarnation_for=incarnations.get, + ) + planner._mixed_current_snapshot_meta = dict( + buffer=planner._host_buffer, + event=planner._copy_event, + owners=slots.copy(), + staged_iteration=9, + staging_sequence=5, + confidence_stamp=planner._host_confidence_stamp, + producer_attempts={r.py_request_id: (r, i, i + 1) for i, r in enumerate(generations)}, + ) + if stale and rank == 1: + incarnations[generations[0].py_request_id] += 100 + runner = SimpleNamespace( + _config=config, + mapping=SimpleNamespace(tp_size=8, pp_size=1, cp_size=1), + _dspark_confidence_enabled=True, + _dspark_trims_submitted_tokens=True, + use_beam_search=False, + _dspark_device_budget=10, + cuda_graph_runner=SimpleNamespace( + agreed_ragged_bucket=192, + ragged_pad_verify_len=1, + ragged_zero_real_high_rows=0, + ), + breakable_cuda_graph_runner=SimpleNamespace( + is_capturing=False, is_warming_up=False, has_graph=lambda _bucket: True + ), + ) + group.append((runtime.MixedRuntime(runner, costs), batch, worker, planner, incarnations)) + votes = [ + mixed.local_vote(batch, worker, planner, 10, rank, executor_eligible=True) + for rank, (mixed, batch, worker, planner, _) in enumerate(group) + ] + if not all_decode: + index = runtime.BODY_BUCKETS.index(192) + geometry = MixedGeometry( + tuple(tuple(v[14:26]) for v in votes), + 192, + tuple(v[28 + index * 3 + 2] for v in votes), + tuple(v[27] for v in votes), + ) + costs = replace(costs, cells=(MixedCostCell(geometry, 10.0, 5.0, "4" * 64),)) + for mixed, *_ in group: + mixed.costs = costs + return group, votes + + +def test_host_mixed_shortens_decode_peer_without_mutating_context(monkeypatch): + group, votes = _group(monkeypatch) + context = copy.deepcopy(group[0][1].context_requests[0].__dict__) + for mixed, batch, worker, planner, _ in group: + assert mixed.apply(batch, worker, planner, votes, 10) + assert mixed.runner._dspark_host_window_step + assert mixed.runner.cuda_graph_runner._dspark_host_window_batch is batch + assert mixed.runner._dspark_device_budget is None + assert any(r.py_verify_len < 3 for r in batch.generation_requests) + assert group[0][1].context_requests[0].__dict__ == context + assert sum(r.py_verify_len + 1 for r in group[1][1].generation_requests) == 192 + + +@pytest.mark.parametrize("failure", ["uncovered", "unready", "identity"]) +def test_missing_mixed_evidence_stays_native_without_disabling_decode(monkeypatch, failure): + group, votes = _group(monkeypatch) + if failure == "uncovered": + for mixed, *_ in group: + mixed.costs = replace(mixed.costs, cells=()) + elif failure == "unready": + mixed, batch, worker, planner, _ = group[1] + planner._copy_event.ready = False + votes[1] = mixed.local_vote(batch, worker, planner, 10, 1, executor_eligible=True) + assert votes[1][5] == 0 + else: + votes[2][6] ^= 1 + for mixed, batch, worker, planner, _ in group: + assert mixed.apply(batch, worker, planner, votes, 10) + assert all(r.py_verify_len == 3 for r in batch.generation_requests) + assert mixed.runner._dspark_confidence_enabled + + +def test_stale_owner_keeps_full_width_without_discarding_other_owners(monkeypatch): + group, votes = _group(monkeypatch, stale=True) + mixed, batch, worker, planner, _ = group[1] + assert votes[1][22] == 63 + mixed.apply(batch, worker, planner, votes, 10) + assert batch.generation_requests[0].py_verify_len == 3 + assert any(r.py_verify_len < 3 for r in batch.generation_requests[1:]) + + +@pytest.mark.parametrize("change", ["incarnation", "context", "first_chunk"]) +def test_ticket_rechecked_before_publishing_lengths(monkeypatch, change): + group, votes = _group(monkeypatch) + rank = 1 if change == "incarnation" else 0 + mixed, batch, worker, planner, incarnations = group[rank] + if change == "incarnation": + incarnations[batch.generation_requests[0].py_request_id] += 1 + elif change == "context": + batch.context_requests[0].context_current_position += 1 + else: + batch.context_requests[0].is_first_context_chunk = True + with pytest.raises(ValueError, match="changed"): + mixed.apply(batch, worker, planner, votes, 10) + assert all(r.py_verify_len is None for r in batch.generation_requests) + + +@pytest.mark.parametrize("unresolved", ["first_chunk", "token_budget"]) +def test_unfinalized_geometry_declines_group_before_compact_publication(monkeypatch, unresolved): + group, votes = _group(monkeypatch) + mixed, batch, worker, planner, _ = group[0] + if unresolved == "first_chunk": + batch.context_requests[0].is_first_context_chunk = True + else: + mixed.runner._config.max_num_tokens = 319 + votes[0] = mixed.local_vote(batch, worker, planner, 10, 0, executor_eligible=True) + assert votes[0][5] == 0 + assert all(r.py_verify_len is None for r in batch.generation_requests) + for mixed, batch, worker, planner, _ in group: + assert mixed.apply(batch, worker, planner, votes, 10) + assert all(r.py_verify_len == 3 for r in batch.generation_requests) + assert mixed.runner._dspark_host_window_step + + +def test_first_chunk_stays_native_across_real_python_resource_preparation(monkeypatch): + group, votes = _group(monkeypatch) + mixed, batch, worker, planner, _ = group[0] + + class _NativeRequestModel(SimpleNamespace): + # Models the native first-chunk predicate; no C++ execution is claimed. + @property + def is_first_context_chunk(self): + return self.context_current_position == self.prepopulated_prompt_len + + values = batch.context_requests[0].__dict__.copy() + values.pop("is_first_context_chunk") + context = _NativeRequestModel(**values) + context.context_chunk_size, context.context_current_position = 256, 0 + context.prepopulated_prompt_len = 0 + batch.context_requests[:] = [context] + batch.context_requests_last_chunk[:] = [context] + votes[0] = mixed.local_vote(batch, worker, planner, 10, 0, executor_eligible=True) + assert votes[0][5] == 0 + for mixed, member, worker, planner, _ in group: + mixed.apply(member, worker, planner, votes, 10) + assert all(r.py_verify_len == 3 for r in member.generation_requests) + phases = [] + + def model_prefix_cache_setter(_batch): + phases.append("prepare") + context.prepopulated_prompt_len = context.context_current_position = 128 + context.context_chunk_size = 128 + + manager = SimpleNamespace( + resource_managers={ + ResourceManagerType.KV_CACHE_MANAGER: SimpleNamespace( + prepare_resources=model_prefix_cache_setter, + report_batch_to_connector=lambda _batch: phases.append("report"), + ) + }, + maybe_fit_token_budget=lambda _batch: phases.append("trim"), + ) + ResourceManager.prepare_resources(manager, batch) + assert phases == ["prepare", "trim", "report"] + assert context.context_chunk_size == 128 + assert all(r.py_verify_len == 3 for r in batch.generation_requests) + + +def test_all_decode_keeps_original_policy_and_device_mode(monkeypatch): + group, votes = _group(monkeypatch, all_decode=True) + for mixed, batch, worker, planner, _ in group: + before = [r.__dict__.copy() for r in batch.generation_requests] + assert not mixed.apply(batch, worker, planner, votes, 10) + assert [r.__dict__ for r in batch.generation_requests] == before + assert planner.device_windows + + +def test_optional_suffix_preserves_legacy_payload(monkeypatch): + group, votes = _group(monkeypatch) + mixed = group[0][0] + legacy, suffix = mixed.split_votes([[1, 2, 3] + vote for vote in votes], 3) + assert legacy == [[1, 2, 3]] * 8 + assert suffix == votes + with pytest.raises(ValueError, match="shape differs"): + mixed.split_votes([[1, 2, 3] + vote[:-1] for vote in votes], 3) + + +def test_unsupported_backend_declines_before_reading_snapshot(monkeypatch): + group, _ = _group(monkeypatch) + mixed, batch, worker, planner, _ = group[1] + mixed.runner._config.prefill_cuda_graph_backend = SimpleNamespace(name="PIECEWISE") + planner._mixed_current_snapshot_meta = None + vote = mixed.local_vote(batch, worker, planner, 10, 1, executor_eligible=True) + assert len(vote) == runtime.WIRE_WORDS + assert vote[5] == 0 + + +def test_host_metadata_skips_device_prologue_and_local_full_graph(): + runner = SimpleNamespace( + _dspark_trims_submitted_tokens=True, + _dspark_device_windows=True, + cuda_graph_runner=SimpleNamespace(agreed_ragged_bucket=192), + get_runtime_tokens_per_gen_step=lambda k: k + 1, + _ragged_token_lens=DecoderRunner._ragged_token_lens, + ) + metadata = SimpleNamespace( + runtime_tokens_per_gen_step=None, ragged_verify_lens=None, device_windows_mode=False + ) + requests = [SimpleNamespace(py_verify_len=1), SimpleNamespace(py_verify_len=3)] + DecoderRunner._publish_gen_token_layout( + runner, metadata, requests, runtime_draft_len=3, host_authoritative=True + ) + assert not metadata.device_windows_mode + assert metadata.ragged_verify_lens == [2, 4] + runner._dspark_host_window_step, runner._dspark_device_budget = True, 5 + assert not DecoderRunner._apply_device_window_prologue(runner, {}, None) + assert runner._dspark_device_budget is None + batch = SimpleNamespace(can_run_cuda_graph=True) + graph = SimpleNamespace( + _dspark_host_window_batch=batch, + is_encoder_decoder=False, + enable_encoder_decoder_mixed_cuda_graph=False, + ) + assert not CUDAGraphRunner._can_run_cuda_graph_batch(graph, batch) + + +def test_producer_stamp_copy_uses_same_event_and_validates_before_rotation(monkeypatch): + monkeypatch.setattr(torch.cuda, "Event", _Event) + planner = DSparkVerifyPlanner(cfg=DSparkScheduleConfig(3), device_windows=True) + planner.stage_confidence(torch.ones(2, 3), torch.tensor([1, 2], dtype=torch.int32)) + first = (planner._host_buffer, planner._host_confidence_stamp, planner._copy_event) + with pytest.raises(ValueError, match="one per slot"): + planner.stage_confidence(torch.zeros(2, 3), torch.ones(3, dtype=torch.int32)) + assert planner._host_buffer is first[0] + assert planner._host_confidence_stamp is first[1] + planner.stage_confidence(torch.zeros(2, 3), torch.tensor([3, 4], dtype=torch.int32)) + assert planner._prev_buffer is first[0] + assert planner._prev_confidence_stamp is first[1] + assert planner._prev_event is first[2] + assert planner._host_confidence_stamp.tolist() == [3, 4] + + +def test_forward_attempt_alone_is_not_confidence_publication(monkeypatch): + group, _ = _group(monkeypatch) + _, batch, worker, planner, _ = group[1] + runtime.stage_mixed_ownership(planner, worker, 10, batch.generation_requests) + assert planner._mixed_current_snapshot_meta["producer_attempts"] == {} + worker._draft_seq_host += 1 + runtime.stage_mixed_ownership(planner, worker, 11, batch.generation_requests) + assert len(planner._mixed_current_snapshot_meta["producer_attempts"]) == 64 + assert ( + planner._mixed_current_snapshot_meta["confidence_stamp"] is planner._host_confidence_stamp + ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py new file mode 100644 index 000000000000..94628fa590f8 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model_engine_runtime.py @@ -0,0 +1,680 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Focused CPU tests for DSpark decoder-runner ragged runtime state.""" + +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch import utils as torch_utils +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.metadata import ( + DeepseekV4TrtllmAttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import DSAtrtllmAttentionMetadata +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner, DSparkCUDAGraphRunner +from tensorrt_llm._torch.pyexecutor.engine import metadata as metadata_module +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import runner as decoder_module +from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine +from tensorrt_llm._torch.speculative.ragged_helpers import ragged_gather_index_lists + + +def _runner(*, secondary=True): + return SimpleNamespace( + enabled=True, + agreed_ragged_bucket=None, + ragged_pad_verify_len=0, + ragged_zero_real_high_rows=0, + supported_batch_sizes=[4, 16], + secondary_padding_dummy_requests={5: object()} if secondary else {}, + _round_up_batch_size=lambda rows: 4 if rows <= 4 else 16, + will_pad_to=lambda *_args: True, + ) + + +def _engine(runner, buckets): + engine = object.__new__(DecoderRunner) + engine.cuda_graph_runner = runner + engine._config = SimpleNamespace( + spec_config=SimpleNamespace(max_draft_len=5), + max_batch_size=8, + cuda_graph_batch_sizes=[1, 4, 8], + use_mrope=False, + ) + engine._dspark_confidence_enabled = True + engine._dspark_trims_submitted_tokens = True + engine._dspark_device_windows = False + engine._dspark_last_padded_bs = None + engine.ragged_verify_token_buckets = lambda _rows: list(buckets) + return engine + + +def test_pinned_host_ping_pongs_without_private_cuda_events() -> None: + engine = SimpleNamespace( + _pinned_host_cache={}, + _pinned_host_active={}, + ) + + first = DecoderRunner._pinned_host(engine, "rows", [1, 2], torch.long) + second = DecoderRunner._pinned_host(engine, "rows", [3, 4], torch.long) + third = DecoderRunner._pinned_host(engine, "rows", [5, 6], torch.long) + + assert first.data_ptr() == third.data_ptr() + assert first.data_ptr() != second.data_ptr() + assert third.tolist() == [5, 6] + + +@pytest.mark.parametrize( + ("verifier_budget", "scheduled_window", "high_rows"), + [(80, 4, 0), (81, 5, 1), (88, 5, 8)], +) +def test_zero_real_exact_fit_uses_one_scheduled_dummy(verifier_budget, scheduled_window, high_rows): + runner = _runner() + engine = _engine(runner, [48, 64, 80, 81, 88, 96]) + dummy = SimpleNamespace(is_attention_dp_dummy=True, py_verify_len=None) + + bucket = DecoderRunner.fit_ragged_verify_lens( + engine, + [dummy], + [scheduled_window], + peer_stats=[[16, 88, 1], [0, 0, 1]], + exact_shape=(16, verifier_budget, 5), + exact_zero_real=True, + ) + + assert bucket == verifier_budget + assert dummy.py_verify_len == scheduled_window + assert engine._dspark_last_num_real == 0 + assert runner.ragged_pad_verify_len == 4 + assert runner.ragged_zero_real_high_rows == high_rows + + +def test_zero_real_exact_fit_requires_the_secondary_dummy(): + runner = _runner(secondary=False) + engine = _engine(runner, [88]) + dummy = SimpleNamespace(is_attention_dp_dummy=True, py_verify_len=None) + + with pytest.raises(RuntimeError, match="disappeared after"): + DecoderRunner.fit_ragged_verify_lens( + engine, + [dummy], + [5], + peer_stats=[[16, 88, 1], [0, 0, 1]], + exact_shape=(16, 88, 5), + exact_zero_real=True, + ) + + +def test_exact_fit_publishes_the_measured_bucket_and_pad_window(): + runner = _runner() + engine = _engine(runner, [8, 24]) + requests = [SimpleNamespace(py_verify_len=None) for _ in range(2)] + + bucket = DecoderRunner.fit_ragged_verify_lens( + engine, + requests, + [2, 2], + peer_stats=[[2, 6, 1]], + exact_shape=(4, 8, 1), + ) + + assert bucket == 8 + assert [request.py_verify_len for request in requests] == [2, 2] + assert runner.agreed_ragged_bucket == 8 + assert runner.ragged_pad_verify_len == 0 + + +def test_full_k_bucket_preserves_the_native_static_graph(): + runner = _runner() + engine = _engine(runner, [24]) + requests = [SimpleNamespace(py_verify_len=None) for _ in range(2)] + + bucket = DecoderRunner.fit_ragged_verify_lens( + engine, + requests, + [5, 5], + peer_stats=[[2, 12, 1]], + exact_shape=(4, 24, 6), + ) + + assert bucket is None + assert runner.agreed_ragged_bucket is None + assert all(request.py_verify_len is None for request in requests) + + +def _warmup_metadata(position_helper, mask_helper): + ratios = [1, 4] + return SimpleNamespace( + _compress_ratios_sorted=ratios, + max_draft_tokens=5, + past_kv_lens_cuda={ratio: torch.empty(8, dtype=torch.int32) for ratio in ratios}, + cu_new_comp_kv_cuda={ratio: torch.empty(9, dtype=torch.int32) for ratio in ratios}, + new_comp_kv_lens_cuda={ratio: torch.empty(8, dtype=torch.int32) for ratio in ratios}, + compressed_position_ids_cuda={ + ratio: torch.empty(48, dtype=torch.int32) for ratio in ratios + }, + compressed_mask_cuda={ratio: torch.empty(48, dtype=torch.bool) for ratio in ratios}, + batch_size_cuda=torch.tensor([17], dtype=torch.int32), + _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(new_comp, cu_new_comp, compressed_mask, batch_size, compress_ratios): + mask_calls.append((torch.is_inference_mode_enabled(), batch_size.tolist(), compress_ratios)) + + metadata = _warmup_metadata(record_positions, record_mask) + engine = SimpleNamespace( + _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, 4])] + assert metadata.batch_size_cuda.tolist() == [17] + for ratio, token_count in ((1, 42), (4, 14)): + assert not metadata.past_kv_lens_cuda[ratio][:7].any() + assert not metadata.new_comp_kv_lens_cuda[ratio][:7].any() + 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() + assert metadata.batch_size_cuda.tolist() == [17] + + +@pytest.mark.parametrize("fail_after_mask", [False, True]) +def test_ragged_compressor_warmup_calls_current_metadata_helpers(monkeypatch, fail_after_mask): + """Run the real five-argument mask and position helper bodies on CPU tensors.""" + actual_mask = DeepseekV4TrtllmAttentionMetadata._compute_compressed_mask + + def mask(new_comp, cu_new_comp, compressed_mask, batch_size, compress_ratios): + actual_mask(new_comp, cu_new_comp, compressed_mask, batch_size, compress_ratios) + assert batch_size is metadata.batch_size_cuda + assert batch_size.tolist() == [7] + assert compressed_mask[1][:42].all() + assert compressed_mask[4][:14].all() + assert metadata.compressed_position_ids_cuda[1][:42].tolist() == list(range(6)) * 7 + assert metadata.compressed_position_ids_cuda[4][:14].tolist() == [0, 4] * 7 + if fail_after_mask: + raise RuntimeError("mask cleanup probe") + + metadata = _warmup_metadata( + DeepseekV4TrtllmAttentionMetadata._compute_gen_compressed_position_ids, mask + ) + original_storage = metadata.batch_size_cuda.data_ptr() + engine = SimpleNamespace( + _dspark_confidence_enabled=True, + _dspark_trims_submitted_tokens=True, + attn_metadata=metadata, + _config=SimpleNamespace(max_batch_size=8, cuda_graph_batch_sizes=[1, 4, 8]), + ) + # Exercise the helpers' existing eager branch, not a fake signature or GPU compile. + monkeypatch.setattr(torch_utils, "is_piecewise_running", lambda: True) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + if fail_after_mask: + with pytest.raises(RuntimeError, match="mask cleanup probe"): + DecoderRunner._warmup_dspark_ragged_compressor_metadata(engine) + else: + DecoderRunner._warmup_dspark_ragged_compressor_metadata(engine) + assert metadata.batch_size_cuda.data_ptr() == original_storage + assert metadata.batch_size_cuda.tolist() == [17] + for ratio, token_count in ((1, 42), (4, 14)): + assert not metadata.past_kv_lens_cuda[ratio][:7].any() + assert not metadata.new_comp_kv_lens_cuda[ratio][:7].any() + assert not metadata.cu_new_comp_kv_cuda[ratio][:8].any() + assert not metadata.compressed_position_ids_cuda[ratio][:token_count].any() + assert not metadata.compressed_mask_cuda[ratio][:token_count].any() + + +def test_ragged_compressor_warmup_requires_trimmed_tokens(): + engine = SimpleNamespace( + _dspark_confidence_enabled=True, + _dspark_trims_submitted_tokens=False, + ) + + DecoderRunner._warmup_dspark_ragged_compressor_metadata(engine) + + +def test_ragged_compressor_warmup_rejects_missing_metadata_contract(): + engine = SimpleNamespace( + _dspark_confidence_enabled=True, + _dspark_trims_submitted_tokens=True, + attn_metadata=SimpleNamespace(), + ) + + with pytest.raises(RuntimeError, match="missing .*max_draft_tokens"): + DecoderRunner._warmup_dspark_ragged_compressor_metadata(engine) + + +@pytest.mark.parametrize("runtime_draft_len", [0, 2, 5]) +def test_publish_gen_token_layout_uses_call_draft_length(runtime_draft_len: int) -> None: + engine = _engine(_runner(), []) + engine.get_runtime_tokens_per_gen_step = lambda draft_len: draft_len + 1 + requests = [ + SimpleNamespace(py_verify_len=min(1, runtime_draft_len)), + SimpleNamespace(py_verify_len=min(3, runtime_draft_len)), + ] + metadata = SimpleNamespace( + runtime_tokens_per_gen_step=0, + ragged_verify_lens=None, + device_windows_mode=False, + ) + assert not hasattr(engine, "runtime_draft_len") + + engine._publish_gen_token_layout(metadata, requests, runtime_draft_len=runtime_draft_len) + + assert metadata.runtime_tokens_per_gen_step == runtime_draft_len + 1 + assert metadata.ragged_verify_lens == [request.py_verify_len + 1 for request in requests] + assert not metadata.device_windows_mode + assert not hasattr(engine, "runtime_draft_len") + + +def test_attach_ragged_layout_refreshes_persistent_buffers_in_place() -> None: + runner = _engine(_runner(), []) + runner._pinned_host_cache = {} + runner._pinned_host_active = {} + runner.ragged_verify_lens_cuda = torch.empty(9, dtype=torch.int32) + runner.ragged_qo_indptr_cuda = torch.empty(10, dtype=torch.int32) + attention = SimpleNamespace(ragged_verify_lens=None) + speculation = SimpleNamespace() + requests = [SimpleNamespace(py_verify_len=1), SimpleNamespace(py_verify_len=3)] + + runner._attach_ragged_verify_layout(speculation, attention, requests) + first_lens_ptr = speculation.verify_lens.data_ptr() + first_prefix_ptr = speculation.qo_indptr.data_ptr() + + assert speculation.verify_lens.tolist() == [2, 4] + assert speculation.qo_indptr.tolist() == [0, 2, 6] + assert speculation.total_verify_tokens == 6 + assert attention.ragged_verify_lens == [2, 4] + + requests[0].py_verify_len = 0 + requests[1].py_verify_len = 2 + runner._attach_ragged_verify_layout(speculation, attention, requests) + + assert speculation.verify_lens.tolist() == [1, 3] + assert speculation.qo_indptr.tolist() == [0, 1, 4] + assert speculation.total_verify_tokens == 4 + assert attention.ragged_verify_lens == [1, 3] + assert speculation.verify_lens.data_ptr() == first_lens_ptr + assert speculation.qo_indptr.data_ptr() == first_prefix_ptr + + +@pytest.mark.parametrize("trim_enabled", [False, True]) +def test_missing_or_disabled_windows_clear_stale_ragged_views(trim_enabled: bool) -> None: + runner = _engine(_runner(), []) + runner._dspark_trims_submitted_tokens = trim_enabled + attention = SimpleNamespace(ragged_verify_lens=[2, 4]) + speculation = SimpleNamespace( + verify_lens=torch.tensor([2, 4]), + qo_indptr=torch.tensor([0, 2, 6]), + total_verify_tokens=6, + ) + + runner._attach_ragged_verify_layout( + speculation, + attention, + [SimpleNamespace(py_verify_len=1), SimpleNamespace(py_verify_len=None)], + ) + + assert speculation.verify_lens is None + assert speculation.qo_indptr is None + assert speculation.total_verify_tokens is None + assert attention.ragged_verify_lens is None + + +def test_model_engine_delegates_ragged_state_to_decoder_owner() -> None: + runner = _engine(_runner(), []) + runner._dspark_sps_cost_table = object() + runner._dspark_exact_candidate_cells = ((4, 12),) + runner._dspark_exact_identity_words = tuple(range(8)) + runner.fit_ragged_verify_lens = lambda *args, **kwargs: (args, kwargs) + worker = object() + runner._get_spec_worker = lambda: worker + engine = object.__new__(PyTorchModelEngine) + engine._cleanup_done = True + engine._runner = runner + + assert engine._dspark_confidence_enabled + assert engine._dspark_trims_submitted_tokens + assert engine._dspark_sps_cost_table is runner._dspark_sps_cost_table + assert engine._dspark_exact_candidate_cells == ((4, 12),) + assert engine._dspark_exact_identity_words == tuple(range(8)) + assert engine._dspark_device_budget is None + assert engine._get_spec_worker() is worker + engine._dspark_device_budget = 7 + assert runner._dspark_device_budget == 7 + assert engine.fit_ragged_verify_lens([1], exact_shape=(4, 12, 3)) == ( + ([1],), + {"exact_shape": (4, 12, 3)}, + ) + + +def test_model_engine_non_decoder_keeps_confidence_inactive() -> None: + engine = object.__new__(PyTorchModelEngine) + engine._cleanup_done = True + engine._runner = SimpleNamespace() + + assert not engine._dspark_confidence_enabled + assert not engine._dspark_trims_submitted_tokens + assert engine._dspark_sps_cost_table is None + assert engine._dspark_exact_candidate_cells == () + assert engine._dspark_exact_identity_words == (0,) * 8 + assert engine._dspark_device_budget is None + assert engine._get_spec_worker() is None + + +def _metadata_inputs(metadata_type): + model = SimpleNamespace( + pretrained_config=SimpleNamespace(num_attention_heads=4, num_key_value_heads=1), + sparse_attention_config=None, + enable_flash_mla=False, + ) + kwargs = dict( + max_batch_size=2, + max_num_tokens=8, + max_beam_width=1, + attention_backend=SimpleNamespace(Metadata=metadata_type), + attention_runtime_features=SimpleNamespace(cache_reuse=False, chunked_prefill=False), + mapping=object(), + cache_indirection=None, + enable_context_mla_with_cached_kv=False, + ) + return model, kwargs + + +@pytest.mark.parametrize("explicit_false", [False, True]) +def test_metadata_helper_keeps_native_constructor_kwargs(explicit_false): + seen = [] + + class NativeMetadata: + def __init__(self, **kwargs): + seen.append(kwargs) + + model, kwargs = _metadata_inputs(NativeMetadata) + if explicit_false: + kwargs["enable_ragged_verification"] = False + metadata_module.build_attention_metadata(model, **kwargs) + + assert "enable_ragged_verification" not in seen[0] + assert seen[0]["num_heads_per_kv"] == 4 + + +def test_metadata_helper_enables_only_supported_ragged_metadata(): + class RaggedMetadata(DSAtrtllmAttentionMetadata): + def __init__(self, **kwargs): + self.constructor_flag = kwargs.pop("enable_ragged_verification", False) + self.enable_ragged_verification = self.constructor_flag + self.ragged_verify_lens = [1, 2] + + model, kwargs = _metadata_inputs(RaggedMetadata) + result = metadata_module.build_attention_metadata( + model, enable_ragged_verification=True, **kwargs + ) + assert result.constructor_flag is True + result._validate_ragged_verification() + + kwargs["attention_backend"] = SimpleNamespace(Metadata=object) + with pytest.raises(ValueError, match="DSA target metadata"): + metadata_module.build_attention_metadata(model, enable_ragged_verification=True, **kwargs) + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_decoder_enables_ragged_metadata_during_construction(monkeypatch, enabled): + seen = [] + + class MetadataProbe(DSAtrtllmAttentionMetadata): + def __init__(self, **kwargs): + self.enable_ragged_verification = kwargs.pop("enable_ragged_verification", False) + self.kv_cache_manager = kwargs["kv_cache_manager"] + self.ragged_verify_lens = None + seen.append(self.enable_ragged_verification) + + model, kwargs = _metadata_inputs(MetadataProbe) + decoder = object.__new__(DecoderRunner) + decoder.attn_metadata = None + decoder._dspark_confidence_enabled = enabled + decoder.model = SimpleNamespace(model_config=model) + decoder._config = SimpleNamespace(**kwargs) + decoder.mapping = kwargs["mapping"] + decoder.cache_indirection_attention = None + monkeypatch.setattr(decoder_module, "resolve_mamba_metadata_cls", lambda _model: None) + kv = object() + + result = decoder._set_up_attn_metadata(kv) + if enabled: + result.ragged_verify_lens = [1, 2] + result._validate_ragged_verification() + + assert seen == [enabled] + assert decoder._set_up_attn_metadata(kv) is result + assert seen == [enabled] + + +@pytest.mark.parametrize("graph_type", [CUDAGraphRunner, DSparkCUDAGraphRunner]) +def test_current_main_warmup_keeps_serving_timeout_only_for_capture(monkeypatch, graph_type): + graph = graph_type.__new__(graph_type) + graph.clear = lambda: None + graph._capture_allowed = False + graph.padding_dummy_requests = {5: object()} + if graph_type is DSparkCUDAGraphRunner: + graph.secondary_padding_dummy_requests = {5: object()} + graph.ragged_zero_real_high_rows = 2 + graph.adp_shape_agreement = object() + serving_phase = False + calls = [] + + @contextmanager + def serving_timeouts(): + nonlocal serving_phase + serving_phase = True + try: + yield + finally: + serving_phase = False + + def warmup(_resources): + calls.append((graph.is_warmup_only, serving_phase, bool(graph.padding_dummy_requests))) + assert graph._capture_allowed + if not graph.is_warmup_only and graph_type is DSparkCUDAGraphRunner: + assert graph.secondary_padding_dummy_requests == {} + assert graph.ragged_zero_real_high_rows == 0 + assert graph.adp_shape_agreement is None + + monkeypatch.setattr(decoder_module, "moe_comm_serving_timeouts", serving_timeouts) + decoder = SimpleNamespace(cuda_graph_runner=graph, _run_cuda_graph_warmup=warmup) + DecoderRunner._warmup_and_capture_cuda_graphs(decoder, object()) + + assert calls == [(True, False, True), (False, True, False)] + assert not graph._capture_allowed + assert not graph.is_warmup_only + + +@pytest.mark.parametrize("worker", [None, SimpleNamespace(verify_planner=None)]) +def test_device_window_prologue_declines_an_unavailable_planner( + worker: SimpleNamespace | None, +) -> None: + decoder = object.__new__(DecoderRunner) + decoder.model = SimpleNamespace(spec_worker=worker) + decoder.cuda_graph_runner = SimpleNamespace(agreed_ragged_bucket=10) + decoder._dspark_device_budget = 10 + decoder._dspark_prev_covers_batch = True + + assert not decoder._apply_device_window_prologue({}, None) + assert decoder._dspark_device_budget is None + + +def test_confidence_without_an_sps_table_is_eager_only() -> None: + decoder = object.__new__(DecoderRunner) + decoder.cuda_graph_runner = SimpleNamespace(enabled=False) + decoder._dspark_confidence_enabled = True + decoder._dspark_trims_submitted_tokens = True + decoder._dspark_sps_cost_table = None + decoder._config = SimpleNamespace( + spec_config=SimpleNamespace(draft_len_schedule=None, max_draft_len=3) + ) + + decoder._capture_generation_cuda_graphs(None) + with pytest.raises(RuntimeError, match="authenticated exact SPS table"): + decoder._get_graphs_to_capture([1, 4]) + + +@pytest.mark.parametrize( + ("trims", "prefix_window", "previous_windows"), + [(True, 1, [3, 3]), (True, 1, [1, 2]), (False, None, [None, None])], +) +def test_overlap_offsets_follow_the_packed_prefix( + trims: bool, prefix_window: int | None, previous_windows: list[int | None] +) -> None: + """Exercise the full input packer with CPU buffers and metadata-only leaves.""" + decoder = object.__new__(DecoderRunner) + mode = SimpleNamespace( + support_overlap_scheduler=lambda: True, extend_ctx=lambda _backend: False + ) + spec_config = SimpleNamespace(spec_dec_mode=mode, is_linear_tree=True) + decoder._config = SimpleNamespace( + use_mrope=False, + max_num_tokens=32, + max_beam_width=1, + is_spec_decode=True, + disable_overlap_scheduler=False, + spec_config=spec_config, + attention_backend="TRTLLM", + enable_attention_dp=False, + prefill_cuda_graph_backend=decoder_module.PrefillCudaGraphBackend.DISABLED, + prefill_cuda_graph_num_tokens=None, + ) + decoder.mapping = SimpleNamespace(pp_size=1, has_cp_helix=lambda: False) + decoder.dist = None + decoder.guided_decoder = None + decoder.model = SimpleNamespace( + model_config=SimpleNamespace(pretrained_config=SimpleNamespace()) + ) + decoder._dspark_trims_submitted_tokens = trims + decoder._dspark_device_windows = False + decoder._prepare_inputs_fast_path = lambda *_args, **_kwargs: None + decoder._can_use_steady_gen_fast_prepare = lambda *_args: False + decoder._new_extra_inputs = lambda: SimpleNamespace( + add_generation_request=lambda _request: None, build=lambda *_args: {} + ) + decoder._lora = SimpleNamespace(build=lambda *_args, **_kwargs: None) + decoder._should_use_full_generation_page_table = lambda *_args: False + decoder.get_runtime_tokens_per_gen_step = lambda draft_len: draft_len + 1 + decoder.use_beam_search = False + decoder.iter_states = {} + decoder._log_cached_kv_tokens_per_req = False + decoder._pinned_host_cache = {} + decoder._pinned_host_active = {} + decoder._ragged_gather_indices = lambda slots, lengths: tuple( + torch.tensor(indices, dtype=torch.long) + for indices in ragged_gather_index_lists(slots, lengths) + ) + for name in ( + "input_ids_cuda", + "draft_tokens_cuda", + "num_accepted_draft_tokens_cuda", + "previous_pos_id_offsets_cuda", + "previous_kv_lens_offsets_cuda", + "previous_pos_indices_cuda", + "previous_batch_indices_cuda", + "position_ids_cuda", + "gather_ids_cuda", + ): + setattr(decoder, name, torch.full((32,), -99, dtype=torch.int32)) + windows = [prefix_window, *previous_windows] + requests = [ + SimpleNamespace( + py_request_id=index + 100, + py_seq_slot=index, + py_batch_idx=None if index == 0 else index, + py_draft_tokens=[11, 12, 13], + py_verify_len=window, + is_dummy=False, + is_attention_dp_dummy=False, + is_cuda_graph_dummy=False, + py_is_first_draft=False, + py_num_accepted_draft_tokens=1, + max_beam_num_tokens=8, + py_prompt_len=4, + py_num_compressed_tokens=0, + get_last_tokens=lambda _beam: 10, + ) + for index, window in enumerate(windows) + ] + scheduled = SimpleNamespace( + context_requests=[], + generation_requests=requests, + num_context_requests=0, + num_generation_requests=3, + is_attention_dp_phase_idle=False, + attention_dp_phase=None, + ) + state = object.__new__(decoder_module.SampleStateTensorsSpec) + state.new_tokens = torch.arange(12, dtype=torch.int32).reshape(4, 3) + state.new_tokens_lens = torch.tensor([2, 3, 4], dtype=torch.int32) + state.next_draft_tokens = torch.arange(9, dtype=torch.int32).reshape(3, 3) + metadata = SimpleNamespace(is_cuda_graph=False, prepare=lambda: None) + + inputs, gather_ids, draft_len = decoder._prepare_tp_inputs( + scheduled, + None, + metadata, + new_tensors_device=state, + enable_spec_decode=True, + runtime_draft_len=3, + is_dummy=False, + ) + + lengths = [1 + window if trims else 4 for window in windows] + prefix_len, *previous_lengths = lengths + expected = [0] * prefix_len + [3] * previous_lengths[0] + [4] * previous_lengths[1] + assert decoder.previous_pos_id_offsets_cuda[: sum(lengths)].tolist() == expected + assert decoder.previous_pos_id_offsets_cuda[sum(lengths) :].count_nonzero().item() == 0 + assert decoder.previous_kv_lens_offsets_cuda[:3].tolist() == [ + 0, + 3 - previous_lengths[0], + 4 - previous_lengths[1], + ] + assert inputs["input_ids"].numel() == sum(lengths) + assert metadata.seq_lens.tolist() == lengths + assert gather_ids.tolist() == list(range(sum(lengths))) + assert draft_len == 3 diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sampling_admission.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sampling_admission.py new file mode 100644 index 000000000000..962cf4e5a7b6 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sampling_admission.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unsupported confidence/sampler combinations fail before model execution.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm._torch.speculative.dspark import validate_confidence_request +from tensorrt_llm._torch.speculative.interface import SpecMetadata + +pytestmark = pytest.mark.cpu_only + + +@pytest.mark.parametrize("enabled,rejection", [(False, True), (True, False), (True, True)]) +@pytest.mark.parametrize( + "controls,greedy", + [ + ({}, True), + ({"temperature": 0.0}, True), + ({"temperature": 1.0}, False), + ({"temperature": 1.0, "top_k": 1}, True), + ({"min_p": 0.2}, False), + ({"min_p": 1.0}, True), + ], +) +def test_request_admission_matches_actual_draft_sampling_controls( + enabled, rejection, controls, greedy +): + config = SimpleNamespace(enable_confidence_scheduling=enabled, use_rejection_sampling=rejection) + sampling = SimpleNamespace(temperature=None, top_k=None, top_p=None, min_p=None) + sampling.__dict__.update(controls) + request = SimpleNamespace(sampling_config=sampling, py_logits_post_processors=None) + if enabled and rejection and not greedy: + with pytest.raises(ValueError, match="advanced rejection draft sampling"): + validate_confidence_request(config, request) + else: + validate_confidence_request(config, request) + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_per_request_logit_processor_is_rejected_at_confidence_admission(enabled): + config = SimpleNamespace(enable_confidence_scheduling=enabled, use_rejection_sampling=False) + request = SimpleNamespace(sampling_config=None, py_logits_post_processors=[object()]) + if enabled: + with pytest.raises(ValueError, match="per-request logits post-processors"): + validate_confidence_request(config, request) + else: + validate_confidence_request(config, request) + + +@pytest.mark.parametrize( + "enabled,rejection,flags", + [ + (True, True, [True, True]), + (True, True, [True, False]), + (True, True, [False, True]), + (True, True, [False]), + (True, False, [False, True]), + (False, True, [False, True]), + ], +) +def test_host_sampler_guard_is_group_uniform_before_graph_lookup(enabled, rejection, flags): + errors = [] + for local_flag in flags: + runner = object.__new__(DecoderRunner) + runner._dspark_confidence_enabled = enabled + runner.mapping = SimpleNamespace(enable_lm_head_tp_in_adp=False, tp_size=len(flags)) + exchange = Mock(return_value=torch.tensor([[flag] for flag in flags], dtype=torch.int64)) + runner.dist = SimpleNamespace(tp_allgather_int64=exchange) + metadata = object.__new__(SpecMetadata) + metadata.use_rejection_sampling = rejection + metadata.is_all_greedy_sample = local_flag + metadata.group_all_greedy_sample = None + runner._sync_group_all_greedy_sample(metadata) + if enabled and rejection and not all(flags): + with pytest.raises(ValueError, match="advanced rejection draft sampling"): + runner._validate_dspark_draft_sampling(metadata) + errors.append(True) + else: + runner._validate_dspark_draft_sampling(metadata) + errors.append(False) + assert exchange.call_count == int(enabled and rejection and len(flags) > 1) + assert len(set(errors)) == 1 + + +@pytest.mark.parametrize( + "enabled,rejection", [(False, False), (False, True), (True, False), (True, True)] +) +def test_capture_gate_retains_only_supported_advanced_variants(enabled, rejection): + runner = object.__new__(DecoderRunner) + runner._dspark_confidence_enabled = enabled + runner._config = SimpleNamespace(spec_config=SimpleNamespace(use_rejection_sampling=rejection)) + assert runner._supports_advanced_draft_sampling() is not (enabled and rejection) + + +@pytest.mark.parametrize("processors", [False, True]) +def test_invalid_confidence_request_is_excluded_before_activation(processors): + executor = object.__new__(PyExecutor) + executor.model_engine = SimpleNamespace( + _dspark_confidence_enabled=True, + spec_config=SimpleNamespace(enable_confidence_scheduling=True, use_rejection_sampling=True), + ) + request = SimpleNamespace( + sampling_config=SimpleNamespace(temperature=1.0, top_k=None, top_p=None, min_p=None), + py_logits_post_processors=[object()] if processors else None, + ) + already_active = object() + executor.waiting_queue = object() + executor.active_requests = [already_active] + executor._fetch_new_requests = Mock(return_value=[request]) + executor._handle_errors = Mock() + executor._configure_csa2_remote_tail = Mock() + + assert executor._fetch_and_activate_new_requests() == [] + assert executor.active_requests == [already_active] + executor._configure_csa2_remote_tail.assert_not_called() + executor._handle_errors.assert_called_once() + assert executor._handle_errors.call_args.kwargs == { + "requests": [request], + "charge_budget": False, + } + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_feature_off_retains_existing_lm_head_tp_group_sync(enabled): + runner = object.__new__(DecoderRunner) + runner._dspark_confidence_enabled = enabled + runner.mapping = SimpleNamespace(enable_lm_head_tp_in_adp=True, tp_size=2) + exchange = Mock(return_value=torch.tensor([[True], [False]], dtype=torch.int64)) + runner.dist = SimpleNamespace(tp_allgather_int64=exchange) + metadata = object.__new__(SpecMetadata) + metadata.use_rejection_sampling = True + metadata.is_all_greedy_sample = True + metadata.group_all_greedy_sample = None + + runner._sync_group_all_greedy_sample(metadata) + + exchange.assert_called_once_with([True]) + assert not metadata.group_all_greedy_sample + assert not metadata.is_all_greedy_sample diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_schedule_fused.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_schedule_fused.py new file mode 100644 index 000000000000..7cbe2eb1711f --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_schedule_fused.py @@ -0,0 +1,543 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm._torch.speculative.dspark_select as device_select +import tensorrt_llm._torch.speculative.dspark_schedule as dspark_schedule +from tensorrt_llm._torch.speculative.dspark_select import select_windows_device +from tensorrt_llm._torch.speculative.dspark_schedule import ( + DSparkScheduleConfig, + schedule_verify_lens_topk, + schedule_verify_lens_topk_fused_fill, +) +from tensorrt_llm._torch.speculative.ragged_helpers import fill_bucket_device + + +def _legacy_schedule_and_fill( + *, + survival: torch.Tensor, + budget: int, + num_real: int, + pad_len: int, + graph_num_tokens: int, + cfg: DSparkScheduleConfig, +) -> torch.Tensor: + rows = torch.arange(survival.shape[0], device=survival.device) + real_survival = torch.where( + (rows < num_real).unsqueeze(1), + survival, + torch.zeros_like(survival), + ) + scheduled = schedule_verify_lens_topk(survival=real_survival, budget=budget, cfg=cfg) + return fill_bucket_device( + scheduled + 1, + num_real=torch.tensor(num_real, device=survival.device), + graph_num_tokens=graph_num_tokens, + max_verify_len=cfg.resolved_max_verify_len + 1, + pad_fill=pad_len, + ) + + +@pytest.mark.parametrize( + "num_rows,num_real,budget,survival_eps", + [ + (1, 0, 0, 1e-6), + (1, 1, 0, 1e-6), + (2, 1, 1, 1e-6), + (2, 1, 9, 1e-6), + (4, 2, 3, 1e-6), + (4, 4, 16, 1e-6), + (8, 4, 9, 1e-6), + (2, 2, 1, 0.35), + (4, 1, 3, 0.35), + (8, 8, 32, 0.35), + ], +) +def test_cpu_fused_helper_matches_established_schedule_and_fill( + num_rows, num_real, budget, survival_eps +): + cfg = DSparkScheduleConfig( + block_size=5, + min_verify_len=1, + max_verify_len=5, + survival_eps=survival_eps, + ) + generator = torch.Generator().manual_seed( + 10000 * num_rows + 1000 * num_real + 10 * budget + int(survival_eps * 10) + ) + survival = torch.cumprod(torch.rand(num_rows, 5, generator=generator), dim=1) + survival[num_real:] = 0 + pad_len = 1 + token_floor = cfg.min_verify_len + 1 + max_token_len = cfg.resolved_max_verify_len + 1 + pad_tokens = (num_rows - num_real) * pad_len + capacity = min(max(budget, 0), num_real * cfg.schedulable_per_request) + minimum = num_real * token_floor + pad_tokens + maximum = num_real * max_token_len + pad_tokens + graph_num_tokens = min(maximum, minimum + capacity + min(num_real, 2)) + + expected = _legacy_schedule_and_fill( + survival=survival, + budget=budget, + num_real=num_real, + pad_len=pad_len, + graph_num_tokens=graph_num_tokens, + cfg=cfg, + ) + actual = schedule_verify_lens_topk_fused_fill( + survival=survival, + budget=budget, + num_real=num_real, + pad_len=pad_len, + graph_num_tokens=graph_num_tokens, + cfg=cfg, + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("num_real,budget,graph_tokens", [(2, 2, 9), (4, 4, 14)]) +def test_fused_workspace_selection_matches_tensor_and_reuses_storage( + num_real, budget, graph_tokens +): + cfg = DSparkScheduleConfig(block_size=3, min_verify_len=1) + workspace = device_select.DeviceWindowWorkspace.allocate( + max_rows=6, max_tokens=16, device="cpu" + ) + workspace.verify_lens.fill_(-777) + pointers = { + name: getattr(workspace, name).data_ptr() + for name in ("verify_lens", "qo_indptr", "req_idx", "kv_correction") + } + logits = torch.tensor([[10.0, 5.0, 1.0], [-2.0, -3.0, -4.0], [0.5, 0.2, 0.1], [1.0, 2.0, 3.0]]) + for confidence in (logits, logits.flip(0)): + kwargs = dict( + confidence_logits=confidence, + slot_idx=torch.tensor([0, 1, 2, 3]), + num_real=num_real, + budget=budget, + graph_num_tokens=graph_tokens, + cfg=cfg, + pad_len=1, + ) + expected = select_windows_device(**kwargs, use_fused_exact=False) + actual = select_windows_device(**kwargs, use_fused_exact=True, workspace=workspace) + assert expected.workspace is None + assert actual.workspace is workspace + for name, pointer in pointers.items(): + result = getattr(actual, name) + assert result.data_ptr() == pointer + torch.testing.assert_close(result, getattr(expected, name), rtol=0, atol=0) + assert expected.verify_lens.data_ptr() != workspace.verify_lens.data_ptr() + assert workspace.verify_lens[4:].tolist() == [-777, -777] + + +@pytest.mark.parametrize("num_rows", [0, 2]) +def test_fused_helper_preserves_output_address_and_unused_capacity(num_rows): + cfg = DSparkScheduleConfig(block_size=3, min_verify_len=1) + output = torch.full((num_rows + 2,), -777, dtype=torch.int32) + actual = schedule_verify_lens_topk_fused_fill( + survival=torch.ones(num_rows, 3), + budget=0, + num_real=num_rows, + pad_len=1, + graph_num_tokens=2 * num_rows, + cfg=cfg, + out=output, + ) + assert actual.untyped_storage().data_ptr() == output.untyped_storage().data_ptr() + assert actual.tolist() == [2] * num_rows + assert output[num_rows:].tolist() == [-777, -777] + + +@pytest.mark.parametrize( + "output,error,message", + [ + (torch.empty(2, dtype=torch.int32, device="meta"), ValueError, "share survival.device"), + (torch.empty(2, dtype=torch.int64), TypeError, "dtype torch.int32"), + (torch.empty(1, 2, dtype=torch.int32), ValueError, "1-D tensor"), + (torch.empty(1, dtype=torch.int32), ValueError, "at least 2 elements"), + (torch.empty(4, dtype=torch.int32)[::2], ValueError, "contiguous"), + ], +) +def test_fused_helper_rejects_incompatible_output_before_dispatch(output, error, message): + cfg = DSparkScheduleConfig(block_size=3, min_verify_len=1) + with pytest.raises(error, match=message): + schedule_verify_lens_topk_fused_fill( + survival=torch.ones(2, 3), + budget=0, + num_real=2, + pad_len=1, + graph_num_tokens=4, + cfg=cfg, + out=output, + ) + + +def test_round_robin_slack_is_not_promoted_into_confidence_budget(): + cfg = DSparkScheduleConfig( + block_size=5, + min_verify_len=1, + max_verify_len=5, + survival_eps=0.5, + ) + survival = torch.tensor( + [ + [1.0, 0.9, 0.1, 0.1, 0.1], + [1.0, 0.8, 0.49, 0.1, 0.1], + ] + ) + actual = schedule_verify_lens_topk_fused_fill( + survival=survival, + budget=2, + num_real=2, + pad_len=1, + graph_num_tokens=7, + cfg=cfg, + ) + # Budget two selects one position per row, then the already-paid graph + # remainder goes to row zero. Promoting that remainder into top-k would + # instead select row one's 0.49 candidate and produce [3, 4]. + assert actual.tolist() == [4, 3] + + +def test_zero_epsilon_declines_fusion_and_preserves_tensor_semantics(monkeypatch): + cfg = DSparkScheduleConfig( + block_size=5, + min_verify_len=1, + max_verify_len=5, + survival_eps=0.0, + ) + + def _must_not_run(**_kwargs): + raise AssertionError("zero-epsilon selector must use the tensor path") + + monkeypatch.setattr(device_select, "schedule_verify_lens_topk_fused_fill", _must_not_run) + result = select_windows_device( + confidence_logits=torch.zeros(3, 5), + slot_idx=torch.tensor([0, 1, 2]), + num_real=2, + budget=2, + graph_num_tokens=7, + cfg=cfg, + pad_len=1, + use_fused_exact=True, + ) + assert int(result.verify_lens.sum()) == 7 + + +def test_tensor_controls_decline_fusion(monkeypatch): + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1, max_verify_len=5) + + def _must_not_run(**_kwargs): + raise AssertionError("tensor controls must use the capture-safe tensor path") + + monkeypatch.setattr(device_select, "schedule_verify_lens_topk_fused_fill", _must_not_run) + result = select_windows_device( + confidence_logits=torch.zeros(2, 5), + slot_idx=torch.tensor([0, 1]), + num_real=torch.tensor(2, dtype=torch.int64), + budget=torch.tensor(1, dtype=torch.int64), + graph_num_tokens=5, + cfg=cfg, + pad_len=1, + use_fused_exact=True, + ) + assert int(result.verify_lens.sum()) == 5 + + +def test_fused_helper_rejects_infeasible_graph_without_dispatch(): + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1, max_verify_len=5) + with pytest.raises(ValueError, match="cannot realize"): + schedule_verify_lens_topk_fused_fill( + survival=torch.ones(2, 5), + budget=4, + num_real=2, + pad_len=1, + graph_num_tokens=5, + cfg=cfg, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("graph_bs", [16, 32, 64, 128, 192, 256]) +def test_cuda_fused_matches_tensor_oracle_across_all_graph_sizes(graph_bs): + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1, max_verify_len=5) + generator = torch.Generator(device="cuda").manual_seed(347805 + graph_bs) + confidence = torch.rand((graph_bs, 5), generator=generator, device="cuda") + survival = torch.cumprod(confidence, dim=1) + # Include exact ties and non-finite candidates in every graph shape. + survival[:2] = 0.75 + survival[2, 2] = torch.nan + survival[3, 1] = torch.inf + num_real = graph_bs - 3 + survival[num_real:] = 0 + budget = min(2 * num_real, num_real * cfg.schedulable_per_request) + pad_len = 1 + minimum = num_real * (cfg.min_verify_len + 1) + (graph_bs - num_real) + graph_num_tokens = minimum + budget + num_real + + expected = _legacy_schedule_and_fill( + survival=survival, + budget=budget, + num_real=num_real, + pad_len=pad_len, + graph_num_tokens=graph_num_tokens, + cfg=cfg, + ) + actual = schedule_verify_lens_topk_fused_fill( + survival=survival, + budget=budget, + num_real=num_real, + pad_len=pad_len, + graph_num_tokens=graph_num_tokens, + cfg=cfg, + ) + torch.cuda.synchronize() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_model_engine_prewarm_requires_device_windows(): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + + engine = object.__new__(DecoderRunner) + engine._dspark_fused_scheduler_enabled = True + engine._dspark_confidence_enabled = True + engine._dspark_trims_submitted_tokens = True + engine._dspark_device_windows = False + + def _must_not_get_worker(): + raise AssertionError("disabled device-window mode must not prewarm fusion") + + engine._get_spec_worker = _must_not_get_worker + DecoderRunner._warmup_dspark_fused_scheduler(engine) + assert engine._dspark_fused_schedule_ready_sizes == set() + + +def test_model_engine_prewarm_records_only_successful_graph_sizes(monkeypatch): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import runner as decoder_runner + + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1, max_verify_len=5) + planner = SimpleNamespace( + cfg=cfg, + exact_cost_table=SimpleNamespace(tables={16: object(), 32: object()}), + ) + engine = object.__new__(DecoderRunner) + engine._dspark_fused_scheduler_enabled = True + engine._dspark_confidence_enabled = True + engine._dspark_trims_submitted_tokens = True + engine._dspark_device_windows = True + engine._get_spec_worker = lambda: SimpleNamespace(verify_planner=planner) + + original_ones = torch.ones + monkeypatch.setattr( + decoder_runner.torch, + "ones", + lambda shape, dtype=None, device=None: original_ones(shape, dtype=dtype), + ) + monkeypatch.setattr(decoder_runner.torch.cuda, "synchronize", lambda: None) + calls = [] + + def _fake_fused(**kwargs): + calls.append(kwargs["num_real"]) + if kwargs["num_real"] == 32: + raise RuntimeError("synthetic compile failure") + return torch.empty(kwargs["num_real"], dtype=torch.int32) + + monkeypatch.setattr(dspark_schedule, "schedule_verify_lens_topk_fused_fill", _fake_fused) + DecoderRunner._warmup_dspark_fused_scheduler(engine) + assert calls == [16, 32] + assert engine._dspark_fused_schedule_ready_sizes == {16} + + +def test_runtime_failure_retires_only_one_shape_and_retries_once(): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + from tensorrt_llm._torch.speculative.dspark_schedule import DSparkFusedScheduleError + + engine = object.__new__(DecoderRunner) + engine._dspark_fused_schedule_ready_sizes = {16, 32} + calls = [] + + def _select(**kwargs): + calls.append(kwargs["use_fused_exact"]) + if kwargs["use_fused_exact"]: + raise DSparkFusedScheduleError("synthetic launch failure") + return "tensor-result" + + result = DecoderRunner._select_dspark_windows_with_fused_fallback( + engine, + select_fn=_select, + selector_kwargs={}, + padded_bs=16, + ) + assert result == "tensor-result" + assert calls == [True, False] + assert engine._dspark_fused_schedule_ready_sizes == {32} + assert engine._dspark_fused_schedule_failure_counts == {16: 1} + + calls.clear() + result = DecoderRunner._select_dspark_windows_with_fused_fallback( + engine, + select_fn=_select, + selector_kwargs={}, + padded_bs=16, + ) + assert result == "tensor-result" + assert calls == [False] + assert engine._dspark_fused_schedule_failure_counts == {16: 1} + + +def _tensor_prologue_case(): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + + confidence = torch.tensor([[10.0, 10.0, 10.0], [-10.0, -10.0, -10.0]]) + planner = SimpleNamespace( + cfg=DSparkScheduleConfig(block_size=3, min_verify_len=1), + apply_calibration=None, + ) + worker = SimpleNamespace( + verify_planner=planner, + staged_confidence_buffer=lambda: confidence, + batch_slot_view=lambda rows: torch.arange(rows), + verified_draft_seq_cuda=lambda: None, + ) + engine = object.__new__(DecoderRunner) + engine.cuda_graph_runner = SimpleNamespace(agreed_ragged_bucket=6, ragged_pad_verify_len=0) + engine._config = SimpleNamespace(use_mrope=False) + engine._dspark_device_budget = 2 + engine._dspark_prev_covers_batch = True + engine._dspark_last_padded_bs = 2 + engine._dspark_last_num_real = 2 + engine._dspark_fused_schedule_ready_sizes = set() + engine._get_spec_worker = lambda: worker + engine.ragged_verify_lens_cuda = torch.tensor([3, 3], dtype=torch.int32) + engine.ragged_qo_indptr_cuda = torch.tensor([0, 3, 6], dtype=torch.int32) + engine.previous_batch_indices_cuda = torch.tensor([0, 1], dtype=torch.int64) + engine.input_ids_cuda = torch.full((6,), -1, dtype=torch.int32) + engine.position_ids_cuda = torch.tensor([10, 11, 12, 20, 21, 22], dtype=torch.int32) + engine.previous_pos_indices_cuda = torch.full((6,), -1, dtype=torch.int64) + engine.previous_pos_id_offsets_cuda = torch.full((6,), -1, dtype=torch.int32) + engine.previous_kv_lens_offsets_cuda = torch.tensor([-1, -2], dtype=torch.int32) + engine.draft_tokens_cuda = torch.full((4,), -1, dtype=torch.int32) + events = [] + attn_metadata = SimpleNamespace( + num_contexts=0, + kv_lens_cuda=torch.tensor([16, 26], dtype=torch.int32), + apply_device_ragged_layout=lambda lens, owners, correction: events.append( + ("attention", lens.clone(), owners.clone(), correction.clone()) + ), + ) + spec_metadata = SimpleNamespace( + remap_expanded_sampling_params=lambda owners, count: events.append( + ("sampling", owners.clone(), count) + ), + ) + new_tensors = SimpleNamespace( + new_tokens=torch.tensor([[100, 200], [101, 201], [102, 202], [103, 203]]), + new_tokens_lens=torch.tensor([2, 1], dtype=torch.int32), + next_draft_tokens=torch.tensor([[11, 12, 13], [21, 22, 23]]), + ) + return ( + engine, + {"spec_metadata": spec_metadata, "attn_metadata": attn_metadata}, + new_tensors, + events, + ) + + +def test_tensor_prologue_rebuilds_owners_positions_and_both_kv_terms(): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + + engine, inputs, new_tensors, events = _tensor_prologue_case() + lens_address = engine.ragged_verify_lens_cuda.data_ptr() + qo_address = engine.ragged_qo_indptr_cuda.data_ptr() + assert DecoderRunner._apply_device_window_prologue(engine, inputs, new_tensors) + assert engine._dspark_device_budget is None + assert engine.ragged_verify_lens_cuda.tolist() == [4, 2] + assert engine.ragged_qo_indptr_cuda.tolist() == [0, 4, 6] + assert engine.ragged_verify_lens_cuda.data_ptr() == lens_address + assert engine.ragged_qo_indptr_cuda.data_ptr() == qo_address + assert engine.input_ids_cuda.tolist() == [100, 101, 102, 103, 200, 201] + assert engine.position_ids_cuda.tolist() == [10, 11, 12, 13, 20, 21] + assert engine.previous_pos_indices_cuda.tolist() == [0, 0, 0, 0, 1, 1] + assert engine.previous_pos_id_offsets_cuda.tolist() == [2, 2, 2, 2, 1, 1] + assert engine.draft_tokens_cuda.tolist() == [11, 12, 13, 21] + assert inputs["attn_metadata"].kv_lens_cuda.tolist() == [18, 24] + assert engine.previous_kv_lens_offsets_cuda.tolist() == [-2, -1] + assert ( + inputs["attn_metadata"].kv_lens_cuda + engine.previous_kv_lens_offsets_cuda + ).tolist() == [16, 23] + assert [event[0] for event in events] == ["sampling", "attention"] + assert events[0][1].tolist() == [0, 0, 0, 0, 1, 1] + assert events[0][2] == 6 + + +def test_decode_only_prologue_declines_local_prefill_without_mutating_buffers(): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + + engine, inputs, new_tensors, events = _tensor_prologue_case() + inputs["attn_metadata"].num_contexts = 1 + tensors = ( + engine.input_ids_cuda, + engine.position_ids_cuda, + engine.ragged_verify_lens_cuda, + engine.ragged_qo_indptr_cuda, + inputs["attn_metadata"].kv_lens_cuda, + engine.previous_kv_lens_offsets_cuda, + ) + originals = [tensor.clone() for tensor in tensors] + assert not DecoderRunner._apply_device_window_prologue(engine, inputs, new_tensors) + for tensor, original in zip(tensors, originals): + torch.testing.assert_close(tensor, original, rtol=0, atol=0) + assert not events + + +@pytest.mark.parametrize("confidence,fused", [(False, False), (True, False), (True, True)]) +def test_workspace_allocation_follows_decoder_state_initialization(monkeypatch, confidence, fused): + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import runner as decoder_runner + + original_empty, original_zeros = torch.empty, torch.zeros + + def cpu_empty(*args, **kwargs): + kwargs["device"] = "cpu" + return original_empty(*args, **kwargs) + + def cpu_zeros(*args, **kwargs): + kwargs["device"] = "cpu" + return original_zeros(*args, **kwargs) + + monkeypatch.setattr(torch, "empty", cpu_empty) + monkeypatch.setattr(torch, "zeros", cpu_zeros) + monkeypatch.setattr(decoder_runner, "prefer_pinned", lambda: False) + monkeypatch.setattr(decoder_runner, "LoraParamBuilder", lambda **kwargs: SimpleNamespace()) + engine = object.__new__(DecoderRunner) + engine._config = SimpleNamespace( + is_spec_decode=True, + max_draft_loop_tokens=3, + max_batch_size=4, + max_num_tokens=16, + max_beam_width=1, + max_seq_len=16, + use_mrope=False, + attention_backend="TRTLLM", + spec_config=SimpleNamespace( + enable_confidence_scheduling=confidence, + enable_fused_confidence_scheduler=fused, + ), + ) + engine._compute_dynamic_draft_len_mapping = lambda: {} + # Follow the constructor order: buffers exist before mutable state/flags. + engine._allocate_decoder_buffers() + engine._init_decoder_state() + if confidence and fused: + workspace = engine._dspark_device_window_workspace + assert workspace.verify_lens.numel() == 5 + assert workspace.req_idx.numel() == 16 + assert workspace.verify_lens.device == engine.input_ids_cuda.device + else: + assert engine._dspark_device_window_workspace is None diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_select.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_select.py new file mode 100644 index 000000000000..aa371b548346 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_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_schedule import DSparkScheduleConfig +from tensorrt_llm._torch.speculative.dspark_select import ( + DeviceWindowWorkspace, + gather_packed_draft_tokens, + select_windows_device, +) + + +def test_workspace_slices_keep_persistent_storage(): + workspace = DeviceWindowWorkspace.allocate(max_rows=4, max_tokens=16, device="cpu") + owners = [ + workspace.verify_lens, + workspace.qo_indptr, + workspace.req_idx, + workspace.kv_correction, + workspace.past_seen, + ] + addresses = [tensor.data_ptr() for tensor in owners] + workspace.validate(device=torch.device("cpu"), num_rows=2, num_tokens=8) + workspace.validate(device=torch.device("cpu"), num_rows=4, num_tokens=16) + assert [tensor.data_ptr() for tensor in owners] == addresses + + +@pytest.mark.parametrize("num_rows,num_tokens", [(5, 16), (4, 17)]) +def test_workspace_rejects_out_of_capacity_shape(num_rows, num_tokens): + workspace = DeviceWindowWorkspace.allocate(max_rows=4, max_tokens=16, device="cpu") + with pytest.raises(ValueError, match="capacity"): + workspace.validate(device=torch.device("cpu"), num_rows=num_rows, num_tokens=num_tokens) + + +def test_workspace_rejects_wrong_row_extent_dtype(): + workspace = DeviceWindowWorkspace.allocate(max_rows=4, max_tokens=16, device="cpu") + workspace.verify_lens = workspace.verify_lens.to(torch.int64) + with pytest.raises(TypeError, match="verify_lens must have dtype"): + workspace.validate(device=torch.device("cpu"), num_rows=4, num_tokens=16) + + +def test_fresh_confidence_spends_budget_on_the_best_real_row(): + result = select_windows_device( + confidence_logits=torch.tensor( + [[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_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..fe29ae5c95cb --- /dev/null +++ b/tests/unittest/_torch/speculative/test_device_window_sampling_layout.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import types + +import pytest +import torch + +from tensorrt_llm._torch.speculative.interface import _RNG_SLOT_SPAN, SpecMetadata + + +def test_device_window_repacking_follows_fresh_token_owners(): + owners = torch.tensor([0, 0, 1, 2, 2, 2], dtype=torch.long) + meta = types.SimpleNamespace( + request_temperatures=torch.tensor([0.2, 0.7, 1.1]), + temperatures=torch.empty(6), + request_top_ks=torch.tensor([2, 7, 11], dtype=torch.int32), + top_ks=torch.empty(6, dtype=torch.int32), + request_top_ps=torch.tensor([0.3, 0.8, 0.95]), + top_ps=torch.empty(6), + request_min_ps=torch.tensor([0.1, 0.2, 0.4]), + min_ps=torch.empty(6), + request_seeds=torch.tensor([101, 202, 303], dtype=torch.int64), + seeds=torch.empty(6, dtype=torch.int64), + request_offsets=torch.tensor([10, 20, 30], dtype=torch.int64), + offsets=torch.empty(6, dtype=torch.int64), + qo_indptr=torch.tensor([0, 2, 3, 6], dtype=torch.int32), + _sampling_params_signature=[("request",), ("expanded",)], + ) + + SpecMetadata.remap_expanded_sampling_params(meta, owners, owners.numel()) + + assert torch.equal(meta.temperatures, meta.request_temperatures.index_select(0, owners)) + assert torch.equal(meta.top_ks, meta.request_top_ks.index_select(0, owners)) + assert torch.equal(meta.top_ps, meta.request_top_ps.index_select(0, owners)) + assert torch.equal(meta.min_ps, meta.request_min_ps.index_select(0, owners)) + assert torch.equal(meta.seeds, meta.request_seeds.index_select(0, owners)) + row_slots = torch.tensor([0, 1, 0, 0, 1, 2], dtype=torch.int64) + expected_offsets = meta.request_offsets.index_select(0, owners) + row_slots * _RNG_SLOT_SPAN + assert torch.equal(meta.offsets, expected_offsets) + assert meta._sampling_params_signature[0] == ("request",) + assert meta._sampling_params_signature[1] is None + + +def test_device_window_repacking_ignores_unallocated_optional_buffers(): + meta = types.SimpleNamespace( + request_temperatures=None, + temperatures=None, + request_top_ks=None, + top_ks=None, + request_top_ps=None, + top_ps=None, + request_min_ps=None, + min_ps=None, + request_seeds=None, + seeds=None, + request_offsets=None, + offsets=None, + _sampling_params_signature=[None, ("expanded",)], + ) + SpecMetadata.remap_expanded_sampling_params(meta, torch.tensor([0]), 1) + assert meta._sampling_params_signature[1] is None + + +@pytest.mark.parametrize("missing_row_starts", [False, True]) +def test_device_window_repacking_rejects_invalid_captured_buffers(missing_row_starts): + meta = types.SimpleNamespace( + request_temperatures=torch.tensor([0.2]), + temperatures=torch.empty(2 if missing_row_starts else 1), + request_top_ks=None, + top_ks=None, + request_top_ps=None, + top_ps=None, + request_min_ps=None, + min_ps=None, + request_seeds=None, + seeds=None, + request_offsets=torch.tensor([10], dtype=torch.int64) if missing_row_starts else None, + offsets=torch.empty(2, dtype=torch.int64) if missing_row_starts else None, + qo_indptr=None, + _sampling_params_signature=[None, ("expanded",)], + ) + expected_error = "ragged row starts" if missing_row_starts else "exceeds the captured" + with pytest.raises(RuntimeError, match=expected_error): + SpecMetadata.remap_expanded_sampling_params(meta, torch.tensor([0, 0]), 2) + + +def test_device_window_repacking_forces_next_host_expansion(): + values = ((0.2, 2, 0.3, 0.1), (0.7, 7, 0.8, 0.2)) + per_request = [(0.2, 2, 0.3, 0.1, 2), (0.7, 7, 0.8, 0.2, 4)] + meta = types.SimpleNamespace( + request_temperatures=torch.tensor([0.2, 0.7]), + temperatures=torch.empty(6), + request_top_ks=torch.tensor([2, 7], dtype=torch.int32), + top_ks=torch.empty(6, dtype=torch.int32), + request_top_ps=torch.tensor([0.3, 0.8]), + top_ps=torch.empty(6), + request_min_ps=torch.tensor([0.1, 0.2]), + min_ps=torch.empty(6), + request_seeds=None, + seeds=None, + request_offsets=None, + offsets=None, + _sampling_params_signature=[values, (values, (2, 4))], + ) + SpecMetadata.remap_expanded_sampling_params(meta, torch.tensor([0, 1, 1, 1, 0, 0]), 6) + + need_request, need_expanded = SpecMetadata._sampling_params_buffers_need_update( + meta, per_request + ) + assert need_request is False + assert need_expanded is True 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/executor/test_dspark_reproduction_factory.py b/tests/unittest/executor/test_dspark_reproduction_factory.py new file mode 100644 index 000000000000..79320e399e45 --- /dev/null +++ b/tests/unittest/executor/test_dspark_reproduction_factory.py @@ -0,0 +1,397 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run in the installed native runtime on CPU before a public GPU caller. + +Actual TorchLlmArgs/LLM validation and factory signatures are exercised; only +costly model construction is intercepted. These do not prove native GPU math. +""" + +import importlib.util +import inspect +import json +import subprocess +from pathlib import Path +from types import SimpleNamespace + +import pytest +import yaml + +from tensorrt_llm._torch.pyexecutor import ( + dspark_observation_runtime, + dspark_request_profile, + py_executor_creator, +) +from tensorrt_llm._torch.speculative import dspark_reproduction +from tensorrt_llm.executor.base_worker import BaseWorker +from tensorrt_llm.llmapi.llm import LLM +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("mode", ["static", "dynamic", "collect"]) +def test_complete_public_parser_actual_llm_validator_and_internal_factory( + tmp_path, monkeypatch, mode +): + root = tmp_path / "source" + root.mkdir() + subprocess.run(["git", "init", str(root)], check=True, capture_output=True) + subprocess.run( + [ + "git", + "-C", + str(root), + "-c", + "user.name=CPU Control", + "-c", + "user.email=cpu-control@example.invalid", + "commit", + "--allow-empty", + "-m", + "CPU", + ], + check=True, + capture_output=True, + ) + model = tmp_path / "model" + model.mkdir() + (model / "config.json").write_text('{"dspark_block_size": 3}', encoding="utf-8") + (model / "model.safetensors.index.json").write_text( + '{"weight_map": {"mtp.0.weight": "not-a-real-model.safetensors"}}', encoding="utf-8" + ) + config = { + "skip_tokenizer_init": True, + "max_batch_size": 2, + "max_seq_len": 128, + "max_num_tokens": 16, + "disable_overlap_scheduler": False, + "cuda_graph_config": {"batch_sizes": [1, 2], "enable_padding": True}, + "speculative_config": { + "decoding_type": "DSpark", + "max_draft_len": 3, + "speculative_model": str(model), + }, + } + config_path = tmp_path / "runtime.yml" + config_path.write_text(yaml.safe_dump(config), encoding="utf-8") + inputs = tmp_path / "inputs.jsonl" + inputs.write_text( + "\n".join(json.dumps({"prompt_token_ids": [4, 5]}) for _ in range(2)), encoding="utf-8" + ) + cost = tmp_path / "structure-is-not-validated-at-argument-parser.json" + cost.write_text("{}", encoding="utf-8") + script = Path(dspark_reproduction.__file__).parents[2] / "tools/dspark_reproduction.py" + spec = importlib.util.spec_from_file_location("_actual_public_endpoint_cli", script) + cli = importlib.util.module_from_spec(spec) + spec.loader.exec_module(cli) + seen = [] + + def prepare_only(args): + llm_args, rows, contract, _ = dspark_reproduction._prepare(args, 1) + assert isinstance(llm_args, TorchLlmArgs) + assert llm_args.speculative_config.max_draft_len == 3 + assert llm_args.speculative_config.enable_confidence_scheduling == (mode == "dynamic") + # Regression: internal executor_cls MUST fail the genuine public LLM + # boundary, not be smuggled into public kwargs and discovered on a GPU. + public_kwargs = {**config, "model": str(model), "backend": "pytorch"} + LLM._validate_args_for_torch_backend(None, public_kwargs) + with pytest.raises(ValueError, match="executor_cls"): + LLM._validate_args_for_torch_backend(None, {**public_kwargs, "executor_cls": object}) + original = py_executor_creator.create_py_executor + + def capture_factory(**kwargs): + # functools.wraps exposes the inner implementation by default; + # bind the real callable, which owns its internal startup timer. + outer_signature = inspect.signature(original, follow_wrapped=False) + assert "_startup_timer" not in outer_signature.parameters + outer_signature.bind(**kwargs) + assert "executor_cls" not in kwargs + assert kwargs["llm_args"] is llm_args + assert isinstance( + kwargs["runtime_constructors"], py_executor_creator._RuntimeConstructors + ) + seen.append(kwargs) + inner = py_executor_creator._create_py_executor_impl + forwarded = [] + + def capture_model_construction(**arguments): + inspect.signature(inner, follow_wrapped=False).bind(**arguments) + timer = arguments["_startup_timer"] + assert isinstance(timer, py_executor_creator._StartupTimer) + assert timer.name == "executor_creation" and timer.started > 0 + assert { + key: value for key, value in arguments.items() if key != "_startup_timer" + } == { + "llm_args": llm_args, + "checkpoint_dir": kwargs["checkpoint_dir"], + "tokenizer": kwargs["tokenizer"], + "profiling_stage_data": None, + "resource_governor_queue": None, + "runtime_constructors": kwargs["runtime_constructors"], + } + forwarded.append(arguments) + return "CPU constructor boundary only" + + # Execute the genuine decorated wrapper and real host-only timer; + # intercept only the inner expensive model construction boundary. + monkeypatch.setattr( + py_executor_creator, "_create_py_executor_impl", capture_model_construction + ) + result = original(**kwargs) + assert len(forwarded) == 1 + return result + + monkeypatch.setattr(py_executor_creator, "create_py_executor", capture_factory) + assert ( + dspark_observation_runtime.create_observation_executor( + llm_args=llm_args, + checkpoint_dir=str(model), + tokenizer=None, + collect_calibration=mode == "collect", + ) + == "CPU constructor boundary only" + ) + assert len(rows) == 2 and contract["batch_size"] == 2 + + monkeypatch.setattr(dspark_reproduction, "run_native", prepare_only) + args = [ + "run", + "--mode", + mode, + "--model", + str(model), + "--config", + str(config_path), + "--inputs", + str(inputs), + "--source-root", + str(root), + "--physical-k", + "3", + "--batch-size", + "2", + "--output-tokens", + "8", + "--timeout", + "3600", + "--output", + str(tmp_path / "private-output"), + ] + if mode == "dynamic": + args += ["--cost-table", str(cost)] + cli.main(args) + assert len(seen) == 1 + # Keep the original setup (including MPI/device/loader handling), overriding + # only its existing internal factory seam. + assert dspark_reproduction.ObservationWorker.setup_engine is BaseWorker.setup_engine + + +@pytest.mark.cpu_only +def test_offline_head_evaluation_rejects_enabled_policy_before_factory(monkeypatch, tmp_path): + def forbidden(**kwargs): + raise AssertionError("must refuse before model construction") + + monkeypatch.setattr(py_executor_creator, "create_py_executor", forbidden) + from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig + + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"mtp.0.weight": "not-a-real-model.safetensors"}}', encoding="utf-8" + ) + spec = DSparkDecodingConfig( + max_draft_len=3, speculative_model=str(tmp_path), enable_confidence_scheduling=True + ) + with pytest.raises(ValueError, match="native full-K"): + dspark_observation_runtime.create_observation_executor( + llm_args=SimpleNamespace(speculative_config=spec), collect_calibration=True + ) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("resource_governor", [False, True]) +def test_profile_factory_keeps_explicit_offline_constructors( + monkeypatch: pytest.MonkeyPatch, resource_governor: bool +) -> None: + profile = object() + llm_args = SimpleNamespace(enable_resource_governor=resource_governor) + seen = [] + result = object() + + def capture_factory(**kwargs): + inspect.signature(py_executor_creator.create_py_executor, follow_wrapped=False).bind( + **kwargs + ) + seen.append(kwargs) + return result + + monkeypatch.setattr(dspark_request_profile, "create_py_executor", capture_factory) + if resource_governor: + with pytest.raises(ValueError, match="resource-governor"): + dspark_request_profile.create_profile_executor(llm_args=llm_args, profile=profile) + assert not seen + return + + assert ( + dspark_request_profile.create_profile_executor( + llm_args=llm_args, profile=profile, checkpoint_dir="model", tokenizer=None + ) + is result + ) + assert len(seen) == 1 + constructors = seen[0]["runtime_constructors"] + assert isinstance(constructors, py_executor_creator._RuntimeConstructors) + assert constructors.model_engine.func is dspark_request_profile.OwnedDSparkProfileEngine + assert constructors.model_engine.keywords == {"profile": profile} + assert constructors.executor is dspark_request_profile.OwnedDSparkProfileExecutor + assert constructors.kv_cache_creator is dspark_request_profile._ProfileKvCacheCreator + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("active_profile", [False, True], ids=["native", "measured"]) +@pytest.mark.parametrize("position_id_offsets", [None, object()], ids=["default", "offsets"]) +def test_profile_replay_forwards_parent_position_offsets( + monkeypatch, active_profile, position_id_offsets +): + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import DSparkCUDAGraphRunner + from tensorrt_llm._torch.pyexecutor.dspark_profile_engine import _DSparkProfileGraphs + + # A newly added parent keyword must be explicitly supported, with the same + # default, before the expensive model warmup reaches this override. + parent_parameters = inspect.signature(DSparkCUDAGraphRunner.replay).parameters + profile_parameters = inspect.signature(_DSparkProfileGraphs.replay).parameters + parent_keywords = { + name: parameter + for name, parameter in parent_parameters.items() + if parameter.kind is inspect.Parameter.KEYWORD_ONLY + } + assert "position_id_offsets" in parent_keywords + assert parent_keywords.keys() <= profile_parameters.keys() + for name, parameter in parent_keywords.items(): + assert profile_parameters[name].kind is inspect.Parameter.KEYWORD_ONLY + assert profile_parameters[name].default == parameter.default + assert profile_parameters["position_id_offsets"].default is None + + key, output, lease = (64, 3, False), object(), object() + attention, speculative = object(), object() + inputs = { + "attn_metadata": attention, + "spec_metadata": speculative, + "optional_metadata": object(), + } + calls, events = [], [] + + def replay(self, actual_key, current_inputs, *, position_id_offsets=None): + calls.append((self, actual_key, current_inputs, position_id_offsets)) + events.append("replay") + return output + + def stable(metadata, snapshot, *, context): + assert metadata is attention or metadata is speculative + assert context["rank"] == 0 + + monkeypatch.setattr(DSparkCUDAGraphRunner, "replay", replay) + monkeypatch.setattr( + _DSparkProfileGraphs, "_assert_graph_tensor_ptrs_stable", staticmethod(stable) + ) + runner = _DSparkProfileGraphs.__new__(_DSparkProfileGraphs) + runner.active_profile_cell = ( + SimpleNamespace(graph_batch_size=64, verifier_budget=0) if active_profile else None + ) + runner.config = SimpleNamespace(mapping=SimpleNamespace(rank=0)) + runner._profile = SimpleNamespace(max_samples=4, physical_k=3) + runner._profile_records = [] + runner._profile_ptrs = {key: ({}, {})} + runner.graph_metadata = {key: inputs} + runner.active_request_ids = ("request",) + runner._profile_pool = SimpleNamespace( + begin=lambda: events.append("begin") or lease, + end=lambda actual: events.append(("end", actual is lease)), + ) + assert runner.replay(key, inputs, position_id_offsets=position_id_offsets) is output + assert len(calls) == 1 + assert calls[0][0] is runner and calls[0][1] is key + assert calls[0][2] is inputs and calls[0][3] is position_id_offsets + assert events == (["begin", "replay", ("end", True)] if active_profile else ["replay"]) + assert len(runner._profile_records) == int(active_profile) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("enabled", [False, True], ids=["native-collect", "dspark"]) +@pytest.mark.parametrize("position_id_offsets", [None, object()], ids=["default", "offsets"]) +def test_observation_replay_forwards_actual_parent_keyword_contract( + tmp_path, monkeypatch, enabled, position_id_offsets +): + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( + CUDAGraphRunner, + DSparkCUDAGraphRunner, + KeyType, + ) + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner + from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine + from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig + + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"mtp.0.weight": "not-a-real-model.safetensors"}}', encoding="utf-8" + ) + spec = DSparkDecodingConfig( + max_draft_len=3, + speculative_model=str(tmp_path), + enable_confidence_scheduling=enabled, + ) + captured = [] + + def capture_factory(**kwargs): + captured.append(kwargs["runtime_constructors"]) + return "CPU factory boundary only" + + monkeypatch.setattr(py_executor_creator, "create_py_executor", capture_factory) + assert ( + dspark_observation_runtime.observation_factory(collect_calibration=not enabled)( + llm_args=SimpleNamespace(speculative_config=spec), tokenizer=None + ) + == "CPU factory boundary only" + ) + engine_class = captured[0].model_engine + monkeypatch.setattr( + PyTorchModelEngine, + "_initialize_decoder_runner", + lambda self, runner_cls, runner_config=None, **kwargs: runner_cls, + ) + engine = engine_class.__new__(engine_class) + decoder_class = engine._initialize_decoder_runner(DecoderRunner) + base = DSparkCUDAGraphRunner if enabled else CUDAGraphRunner + parent_parameters = inspect.signature(base.replay).parameters + monkeypatch.setattr(base, "__init__", lambda self, config: None) + decoder = decoder_class.__new__(decoder_class) + graph = decoder._create_cuda_graph_runner(SimpleNamespace()) + assert type(graph).__mro__[1] is base + parameters = inspect.signature(type(graph).replay).parameters + keywords = { + name: parameter + for name, parameter in parent_parameters.items() + if parameter.kind is inspect.Parameter.KEYWORD_ONLY + } + assert "position_id_offsets" in keywords + assert keywords.keys() <= parameters.keys() + for name, parameter in keywords.items(): + assert parameters[name].kind is inspect.Parameter.KEYWORD_ONLY + assert parameters[name].default == parameter.default + + key = KeyType(64, 3, False, ragged_verify_bucket=64 if enabled else None) + inputs = {"attn_metadata": object(), "spec_metadata": object(), "optional": object()} + output, calls = object(), [] + + def replay(self, actual_key, actual_inputs, *, position_id_offsets=None): + calls.append((self, actual_key, actual_inputs, position_id_offsets)) + return output + + monkeypatch.setattr(base, "replay", replay) + assert graph.replay(key, inputs, position_id_offsets=position_id_offsets) is output + assert len(calls) == 1 + assert calls[0][0] is graph and calls[0][1] is key + assert calls[0][2] is inputs and calls[0][3] is position_id_offsets + assert graph.observed_routes == { + "native_replays": int(not enabled), + "compact_replays": int(enabled), + } + assert len(graph.observed_keys) == 1 + with pytest.raises(TypeError): + graph.replay(key, inputs, unsupported_next_keyword=object()) diff --git a/tests/unittest/llmapi/test_dspark_workload_admission_config.py b/tests/unittest/llmapi/test_dspark_workload_admission_config.py new file mode 100644 index 000000000000..bf144525a4d0 --- /dev/null +++ b/tests/unittest/llmapi/test_dspark_workload_admission_config.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Startup-boundary tests for DSpark workload admission.""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from tensorrt_llm.llmapi.llm_args import CudaGraphConfig, DSparkDecodingConfig, TorchLlmArgs + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("outcome", ["missing", "negative", "stale"]) +def test_dspark_workload_admission_disables_before_runtime_construction(tmp_path, outcome): + (tmp_path / "config.json").write_text('{"dspark_block_size": 2}') + spec_cfg = DSparkDecodingConfig( + max_draft_len=2, + speculative_model=str(tmp_path), + enable_confidence_scheduling=True, + validate_confidence_artifact_freshness=True, + confidence_sps_table_path=(None if outcome == "missing" else "/sealed/sps.json"), + confidence_sps_live_fingerprint_path=( + None if outcome == "missing" else "/sealed/live.json" + ), + confidence_admission_receipt_path=( + None if outcome == "missing" else "/sealed/admission.json" + ), + confidence_admission_receipt_sha256=(None if outcome == "missing" else "a" * 64), + ) + if outcome == "negative": + patch_kwargs = {"return_value": SimpleNamespace(admitted=False)} + elif outcome == "stale": + patch_kwargs = {"side_effect": ValueError("stale receipt")} + else: + patch_kwargs = {"side_effect": TypeError("missing receipt")} + with patch( + "tensorrt_llm._torch.speculative.dspark_planner.load_confidence_workload_admission", + **patch_kwargs, + ): + # No CUDA-graph configuration is supplied. Construction succeeds only + # if admission resolves the feature off before the confidence-specific + # environment validator observes it. + args = TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + ) + + assert args.speculative_config.max_draft_len == 2 + assert args.speculative_config.enable_confidence_scheduling is False + assert args.speculative_config.confidence_sps_table_path is None + assert args.speculative_config.confidence_sps_live_fingerprint_path is None + assert args.speculative_config.confidence_admission_receipt_path is None + assert args.speculative_config.confidence_admission_receipt_sha256 is None + + +@pytest.mark.cpu_only +def test_dspark_positive_workload_admission_preserves_confidence_path(tmp_path): + (tmp_path / "config.json").write_text('{"dspark_block_size": 2}') + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"mtp.0.weight": "model.safetensors"}}' + ) + spec_cfg = DSparkDecodingConfig( + max_draft_len=2, + speculative_model=str(tmp_path), + enable_confidence_scheduling=True, + validate_confidence_artifact_freshness=True, + confidence_sps_table_path="/sealed/sps.json", + confidence_sps_live_fingerprint_path="/sealed/live.json", + confidence_admission_receipt_path="/sealed/admission.json", + confidence_admission_receipt_sha256="a" * 64, + ) + with patch( + "tensorrt_llm._torch.speculative.dspark_planner.load_confidence_workload_admission", + return_value=SimpleNamespace(admitted=True), + ): + args = TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + cuda_graph_config=CudaGraphConfig(batch_sizes=[1], enable_padding=True), + ) + + assert args.speculative_config.enable_confidence_scheduling is True + assert args.speculative_config.confidence_admission_receipt_path == "/sealed/admission.json" + assert args.speculative_config.confidence_admission_receipt_sha256 == "a" * 64 + + +@pytest.mark.cpu_only +def test_dspark_reuse_keeps_enabled_without_workload_admission(tmp_path): + (tmp_path / "config.json").write_text('{"dspark_block_size": 2}') + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"mtp.0.weight": "model.safetensors"}}' + ) + spec_cfg = DSparkDecodingConfig( + max_draft_len=2, + speculative_model=str(tmp_path), + enable_confidence_scheduling=True, + confidence_sps_table_path="/original/sps.json", + confidence_sts_path="/original/sts.json", + ) + with patch( + "tensorrt_llm._torch.speculative.dspark_planner.load_confidence_workload_admission", + side_effect=AssertionError("Reuse must not construct an admission receipt"), + ): + args = TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + cuda_graph_config=CudaGraphConfig(batch_sizes=[1], enable_padding=True), + ) + + assert args.speculative_config.enable_confidence_scheduling is True + assert args.speculative_config.validate_confidence_artifact_freshness is False + assert args.speculative_config.confidence_admission_receipt_path is None + assert args.speculative_config.confidence_admission_receipt_sha256 is None + assert args.speculative_config.confidence_sps_table_path == "/original/sps.json" + assert args.speculative_config.confidence_sts_path == "/original/sts.json" + + +@pytest.mark.cpu_only +def test_dspark_reuse_preserves_padded_graph_requirement(tmp_path): + (tmp_path / "config.json").write_text('{"dspark_block_size": 2}') + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"mtp.0.weight": "model.safetensors"}}' + ) + spec_cfg = DSparkDecodingConfig( + max_draft_len=2, + speculative_model=str(tmp_path), + enable_confidence_scheduling=True, + confidence_sps_table_path="/original/sps.json", + ) + with pytest.raises(ValueError, match="cuda_graph_config"): + TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + cuda_graph_config=None, + ) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index dd0e9298d4d4..0756a56c066d 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -384,7 +384,7 @@ def test_llm_args_with_kvcache_config(self): llm_args_dict = self._yaml_to_dict(yaml_content) llm_args = TorchLlmArgs(model=llama_model_path, **llm_args_dict) - assert llm_args.kv_cache_config.enable_block_reuse == True + assert llm_args.kv_cache_config.enable_block_reuse assert llm_args.kv_cache_config.max_tokens == 1024 assert llm_args.kv_cache_config.max_attention_window == [ 1024, 1024, 1024 @@ -1198,7 +1198,7 @@ def test_KvCacheConfig_declaration(): attention_dp_events_gather_period_ms=10) pybind_config = config._to_pybind() - assert pybind_config.enable_block_reuse == True + assert pybind_config.enable_block_reuse assert pybind_config.max_tokens == 1024 assert pybind_config.max_attention_window == [1024, 1024, 1024] assert pybind_config.free_gpu_memory_fraction == 0.5 @@ -1234,8 +1234,8 @@ def test_KvCacheConfig_declaration(): ).kv_cache_event_hash_algo == "v2_sha256" assert KvCacheConfig(kv_cache_event_hash_algo="v2_sha256_64" ).kv_cache_event_hash_algo == "v2_sha256_64" - assert pybind_config.enable_partial_reuse == True - assert pybind_config.copy_on_partial_reuse == True + assert pybind_config.enable_partial_reuse + assert pybind_config.copy_on_partial_reuse assert pybind_config.attention_dp_events_gather_period_ms == 10 assert BlockReuseConfig( policy="per_conversation").policy == "per_conversation" @@ -1864,8 +1864,8 @@ def test_DynamicBatchConfig_declaration(): pybind_config = PybindMirror.maybe_to_pybind(config) - assert pybind_config.enable_batch_size_tuning == True - assert pybind_config.enable_max_num_tokens_tuning == True + assert pybind_config.enable_batch_size_tuning + assert pybind_config.enable_max_num_tokens_tuning assert pybind_config.dynamic_batch_moving_average_window == 10 @@ -3112,7 +3112,7 @@ def test_dynamic_batch_config_arbitrary_args(self): config = DynamicBatchConfig(enable_batch_size_tuning=True, enable_max_num_tokens_tuning=True, dynamic_batch_moving_average_window=8) - assert config.enable_batch_size_tuning == True + assert config.enable_batch_size_tuning # Arbitrary arguments should be rejected with pytest.raises( @@ -3157,7 +3157,7 @@ def test_kv_cache_config_arbitrary_args(self): """Test that KvCacheConfig rejects arbitrary arguments.""" # Valid arguments should work config = KvCacheConfig(enable_block_reuse=True, max_tokens=1024) - assert config.enable_block_reuse == True + assert config.enable_block_reuse assert config.max_tokens == 1024 # Arbitrary arguments should be rejected @@ -3172,8 +3172,8 @@ def test_extended_runtime_perf_knob_config_arbitrary_args(self): # Valid arguments should work config = ExtendedRuntimePerfKnobConfig(multi_block_mode=True, cuda_graph_mode=False) - assert config.multi_block_mode == True - assert config.cuda_graph_mode == False + assert config.multi_block_mode + assert not config.cuda_graph_mode # Arbitrary arguments should be rejected with pytest.raises( @@ -3288,8 +3288,8 @@ def test_torch_compile_config_arbitrary_args(self): # Valid arguments should work config = TorchCompileConfig(enable_fullgraph=True, enable_inductor=False) - assert config.enable_fullgraph == True - assert config.enable_inductor == False + assert config.enable_fullgraph + assert not config.enable_inductor # Arbitrary arguments should be rejected with pytest.raises( @@ -3775,8 +3775,8 @@ def test_model_defaults_application(self): assert self.get_model_defaults_called modified_args = llm._executor.engine.model_engine.llm_args - assert modified_args.enable_chunked_prefill == True - assert modified_args.kv_cache_config.enable_block_reuse == False + assert modified_args.enable_chunked_prefill + assert not modified_args.kv_cache_config.enable_block_reuse assert modified_args.kv_cache_config.free_gpu_memory_fraction == 0.75 @pytest.mark.part0 @@ -3796,10 +3796,10 @@ def test_user_overrides_respected(self): assert self.get_model_defaults_called modified_args = llm._executor.engine.model_engine.llm_args - assert modified_args.enable_chunked_prefill == False + assert not modified_args.enable_chunked_prefill assert modified_args.max_batch_size == 42 assert modified_args.max_input_len == 256 - assert modified_args.kv_cache_config.enable_block_reuse == True + assert modified_args.kv_cache_config.enable_block_reuse @pytest.mark.part0 def test_partial_user_override(self): @@ -3816,8 +3816,8 @@ def test_partial_user_override(self): modified_args = llm._executor.engine.model_engine.llm_args assert modified_args.max_batch_size == 42 - assert modified_args.enable_chunked_prefill == True - assert modified_args.kv_cache_config.enable_block_reuse == False + assert modified_args.enable_chunked_prefill + assert not modified_args.kv_cache_config.enable_block_reuse @pytest.mark.part0 def test_empty_nested_config_preserves_defaults(self): @@ -3843,7 +3843,7 @@ def test_empty_nested_config_preserves_defaults(self): # Model defaults set enable_block_reuse=False and # free_gpu_memory_fraction=0.75. An empty KvCacheConfig() # should not prevent these from being applied. - assert modified_args.kv_cache_config.enable_block_reuse == False + assert not modified_args.kv_cache_config.enable_block_reuse assert modified_args.kv_cache_config.free_gpu_memory_fraction == 0.75 @@ -3914,8 +3914,7 @@ def _get_qualified_name(cls: type) -> str: @pytest.mark.cpu_only class TestPydanticBestPractices: - """Ensure that the user-facing LlmArgs and its subfields follow Pydantic best practices. - """ + """Check user-facing LlmArgs and subfields for Pydantic best practices.""" # Fields exempt from Pydantic compatibility checks due to typing limitations or other edge cases. # Avoid adding to this list unless absolutely necessary, especially if a field is user-facing. @@ -5260,3 +5259,154 @@ def test_deepseek_resolves_auto_to_python_on_nixl(self) -> None: cfg = self._pretrained_config(["DeepseekV3ForCausalLM"], "deepseek_v3") _resolve_transceiver_runtime_auto(args, DeepseekV3ForCausalLM, cfg) assert args.cache_transceiver_config.transceiver_runtime == "PYTHON" + + +class TestDSparkConfidenceScheduling: + """Config surface for DSpark confidence-scheduled verification.""" + + def _cfg(self, **kw): + from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig + kw.setdefault("max_draft_len", 5) + return DSparkDecodingConfig(**kw) + + def _scheduled(self, **kw): + """Scheduling on with the required profiled cost table.""" + kw["enable_confidence_scheduling"] = True + kw.setdefault("confidence_sps_table_path", "/tmp/sps.json") + kw.setdefault("confidence_sps_live_fingerprint_path", + "/tmp/runtime.json") + return self._cfg(**kw) + + def test_defaults_are_off(self): + c = self._cfg() + assert c.enable_confidence_scheduling is False + assert c.enable_fused_confidence_scheduler is False + + def test_knobs_require_the_master_switch(self): + for kw in ({ + "confidence_sts_path": "/tmp/sts.json" + }, { + "confidence_sps_table_path": "/tmp/sps.json" + }, { + "confidence_sps_live_fingerprint_path": "/tmp/runtime.json" + }, { + "confidence_admission_receipt_path": "/tmp/admission.json" + }, { + "confidence_admission_receipt_sha256": "a" * 64 + }): + with pytest.raises(ValueError, + match="enable_confidence_scheduling"): + self._cfg(**kw) + + def test_admission_artifacts_are_resolved_at_torch_args_boundary(self): + c = self._cfg(enable_confidence_scheduling=True) + assert c.enable_confidence_scheduling is True + + def test_removed_legacy_verify_len_tiers_are_rejected(self): + with pytest.raises(ValidationError, + match="confidence_verify_len_tiers"): + self._scheduled(confidence_verify_len_tiers=[1, 3, 5]) + + def test_physical_k1_resolves_confidence_scheduling_to_static(self): + c = self._scheduled( + max_draft_len=1, + confidence_sts_path="/tmp/sts.json", + confidence_admission_receipt_path="/tmp/admission.json", + confidence_admission_receipt_sha256="a" * 64, + enable_fused_confidence_scheduler=True, + ) + + assert c.enable_confidence_scheduling is False + assert c.enable_fused_confidence_scheduler is False + assert c.confidence_sts_path is None + assert c.confidence_sps_table_path is None + assert c.confidence_sps_live_fingerprint_path is None + assert c.confidence_admission_receipt_path is None + assert c.confidence_admission_receipt_sha256 is None + + def test_physical_k1_needs_no_confidence_artifacts(self): + c = self._cfg(max_draft_len=1, enable_confidence_scheduling=True) + + assert c.enable_confidence_scheduling is False + assert c.max_total_draft_tokens == 1 + assert c.tokens_per_gen_step == 2 + + def test_physical_k1_remains_valid_when_scheduling_is_off(self): + c = self._cfg(max_draft_len=1) + assert c.enable_confidence_scheduling is False + assert c.max_draft_len == 1 + assert c.max_total_draft_tokens == 1 + assert c.tokens_per_gen_step == 2 + + def test_fused_scheduler_is_independently_gated(self): + with pytest.raises(ValueError, + match="enable_fused_confidence_scheduler"): + self._cfg(enable_fused_confidence_scheduler=True) + c = self._scheduled(enable_fused_confidence_scheduler=True) + assert c.enable_fused_confidence_scheduler is True + + @staticmethod + def _environment(**overrides): + values = dict( + cuda_graph_config=SimpleNamespace(enable_padding=True), + speculative_config=SimpleNamespace( + draft_is_embedded_in_target=True), + guided_decoding_backend=None, + enable_lora=False, + lora_config=None, + sparse_attention_config=None, + pipeline_parallel_size=1, + ) + values.update(overrides) + return SimpleNamespace(**values) + + def test_confidence_environment_rejects_standalone_dspark(self): + environment = self._environment(speculative_config=SimpleNamespace( + draft_is_embedded_in_target=False)) + with pytest.raises(ValueError, + match="only the embedded DeepSeek-V4 draft"): + TorchLlmArgs._validate_dspark_confidence_environment(environment) + + @pytest.mark.parametrize( + ("overrides", "message"), + [ + ({ + "cuda_graph_config": None + }, "requires cuda_graph_config"), + ( + { + "cuda_graph_config": SimpleNamespace(enable_padding=False) + }, + "enable_padding=True", + ), + ({ + "guided_decoding_backend": "xgrammar" + }, "guided decoding"), + ({ + "enable_lora": True + }, "does not yet support LoRA"), + ({ + "pipeline_parallel_size": 2 + }, "does not yet support pipeline parallelism"), + ( + { + "sparse_attention_config": + SimpleNamespace( + enable_heuristic_topk=True, + use_cute_dsl_topk=False, + use_cute_dsl_paged_mqa_logits=False, + ) + }, + "enable_heuristic_topk", + ), + ], + ) + def test_confidence_environment_rejects_unsupported_runtime_combinations( + self, overrides, message): + environment = self._environment(**overrides) + with pytest.raises(ValueError, match=message): + TorchLlmArgs._validate_dspark_confidence_environment(environment) + + def test_confidence_environment_accepts_padded_cuda_graphs(self): + TorchLlmArgs._validate_dspark_confidence_environment( + self._environment()) diff --git a/tests/unittest/tools/test_dspark_calibration.py b/tests/unittest/tools/test_dspark_calibration.py new file mode 100644 index 000000000000..1ab167646704 --- /dev/null +++ b/tests/unittest/tools/test_dspark_calibration.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Actual CPU STS maths, input contracts and standalone CLI controls.""" + +import hashlib +import importlib.util +import json +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +import torch + +REPO = Path(__file__).resolve().parents[3] +SCRIPT = REPO / "tensorrt_llm/tools/dspark_calibration.py" +SPEC = importlib.util.spec_from_file_location("standalone_dspark_calibration", SCRIPT) +calibration = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(calibration) + + +class CalibrationCpuTests(unittest.TestCase): + def test_known_cumulative_prefix_target_and_grid_order(self) -> None: + logits = torch.zeros((4, 2), dtype=torch.float32) + labels = torch.tensor([[1, 1], [1, 0], [0, 0], [0, 0]], dtype=torch.float32) + grid = calibration.default_temperature_grid() + result = calibration.fit_sts_temperatures(logits=logits, prefix_mask=labels, grid=grid) + self.assertEqual(grid.numel(), 41) + self.assertEqual(result["temperatures"], [float(grid[0]), float(grid[0])]) + self.assertEqual(result["ece_before"], [0.0, 0.0]) + self.assertEqual(result["ece_after"], [0.0, 0.0]) + + def test_weighted_ece_not_unweighted_bin_mean(self) -> None: + probs = torch.tensor([0.1, 0.1, 0.1, 0.9]) + targets = torch.tensor([0.0, 0.0, 0.0, 0.0]) + expected = float((probs.to(torch.float64).sum() / 4).item()) + self.assertAlmostEqual( + calibration.expected_calibration_error(probs=probs, targets=targets, num_bins=10), + expected, + ) + + def test_physical_k_is_not_fixed_to_pilot_three(self) -> None: + for width in (1, 2, 3, 4, 5, 8): + with self.subTest(width=width): + logits = torch.zeros((4, width)) + labels = torch.zeros((4, width)) + result = calibration.fit_sts_temperatures( + logits=logits, prefix_mask=labels, grid=calibration.default_temperature_grid() + ) + self.assertEqual(len(result["temperatures"]), width) + self.assertTrue(all(value > 0 for value in result["temperatures"])) + + def test_malformed_rows_and_censored_prefix_order_refused(self) -> None: + grid = calibration.default_temperature_grid() + rows = [ + (torch.empty((0, 3)), torch.empty((0, 3))), + (torch.zeros((2, 3)), torch.zeros((2, 2))), + (torch.zeros((2, 3), dtype=torch.int32), torch.zeros((2, 3))), + (torch.tensor([[float("nan")]]), torch.zeros((1, 1))), + (torch.zeros((1, 1)), torch.tensor([[float("inf")]])), + (torch.zeros((1, 1)), torch.tensor([[0.5]])), + (torch.zeros((1, 3)), torch.tensor([[1, 0, 1]])), + ] + for logits, labels in rows: + with self.subTest(shape=list(logits.shape)), self.assertRaises(ValueError): + calibration.fit_sts_temperatures(logits=logits, prefix_mask=labels, grid=grid) + for bad in (torch.tensor([]), torch.tensor([0.0]), torch.tensor([float("inf")])): + with self.subTest(grid=bad), self.assertRaises(ValueError): + calibration.fit_sts_temperatures( + logits=torch.zeros((2, 1)), prefix_mask=torch.zeros((2, 1)), grid=bad + ) + with self.assertRaises(ValueError): + calibration.fit_sts_temperatures( + logits=torch.zeros((2, 1)), prefix_mask=torch.zeros((2, 1)), grid=grid, num_bins=0 + ) + + def test_safe_shard_pairing_width_and_duplicate_contracts(self) -> None: + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "paired.pt" + logits = torch.zeros((4, 3)) + mask = torch.zeros((4, 3)) + torch.save( + {"meta": {"pairing": "draft_seq_ring"}, "logits": logits, "prefix_mask": mask}, path + ) + actual_logits, actual_mask, hashes = calibration.load_shards([path], 3) + self.assertTrue(torch.equal(logits, actual_logits)) + self.assertTrue(torch.equal(mask, actual_mask)) + self.assertEqual(hashes[str(path)], hashlib.sha256(path.read_bytes()).hexdigest()) + with self.assertRaisesRegex(ValueError, "Duplicate"): + calibration.load_shards([path, path], 3) + with self.assertRaisesRegex(ValueError, "width"): + calibration.load_shards([path], 5) + torch.save( + { + "meta": {"pairing": "unpaired_live_buffer"}, + "logits": logits, + "prefix_mask": mask, + }, + path, + ) + with self.assertRaisesRegex(ValueError, "same stamped draft"): + calibration.load_shards([path], 3) + torch.save({"logits": logits, "prefix_mask": mask}, path) + with self.assertRaisesRegex(ValueError, "pairing metadata"): + calibration.load_shards([path], 3) + + def test_real_cli_from_unrelated_cwd_and_exclusive_output(self) -> None: + with tempfile.TemporaryDirectory() as directory: + base = Path(directory) + shards, output = base / "paired.pt", base / "sts.json" + logits = torch.zeros((4, 3)) + mask = torch.tensor([[1, 1, 1], [1, 1, 0], [1, 0, 0], [0, 0, 0]], dtype=torch.float32) + torch.save( + {"meta": {"pairing": "draft_seq_ring"}, "logits": logits, "prefix_mask": mask}, + shards, + ) + argv = [ + sys.executable, + "-I", + "-B", + str(SCRIPT), + "--shards", + str(shards), + "--physical-k", + "3", + "--output", + str(output), + "--minimum-samples", + "4", + ] + result = subprocess.run(argv, cwd=directory, capture_output=True, text=True, timeout=60) + self.assertEqual(result.returncode, 0, result.stderr) + payload = json.loads(output.read_bytes()) + self.assertEqual(payload["sts_temperatures"], payload["temperatures"]) + self.assertEqual(payload["num_samples"], 4) + self.assertEqual(payload["physical_k"], 3) + self.assertFalse(payload["task_accuracy_or_economic_admission_qualified"]) + self.assertFalse(payload["artifact_freshness_validated"]) + before = output.read_bytes() + again = subprocess.run(argv, cwd=directory, capture_output=True, text=True, timeout=60) + self.assertNotEqual(again.returncode, 0) + self.assertIn("FileExistsError", again.stderr) + self.assertEqual(output.read_bytes(), before) + + def test_no_native_trtllm_import_and_output_does_not_fabricate_admission(self) -> None: + text = SCRIPT.read_text() + self.assertNotIn("import tensorrt_llm", text) + self.assertNotIn("torch.cuda", text) + self.assertIn("weights_only=True", text) + self.assertNotIn('"admitted": True', text) + self.assertNotIn("/scratch/", text) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/unittest/tools/test_dspark_cost_export.py b/tests/unittest/tools/test_dspark_cost_export.py new file mode 100644 index 000000000000..87066313a6ad --- /dev/null +++ b/tests/unittest/tools/test_dspark_cost_export.py @@ -0,0 +1,409 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU artifact controls, not measured GPU costs or performance evidence.""" + +import copy +import hashlib +import importlib.util +import json +import statistics +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[3] +SOURCE = ROOT / "tensorrt_llm/tools/dspark_cost_export.py" +SPEC = importlib.util.spec_from_file_location("public_dspark_cost_export_cpu", SOURCE) +MODULE = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = MODULE +SPEC.loader.exec_module(MODULE) + + +def _digest(value: object) -> str: + return hashlib.sha256( + json.dumps(value, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + + +def _costs(k: int = 3, graphs: tuple[int, ...] = (16, 64), world: int = 8) -> dict: + """Typed synthetic fixtures only; these numbers are not benchmark prices.""" + fingerprint = { + "global_graph_batch_sizes": [g * world for g in graphs], + "gpu": "CPU_SCHEMA_FIXTURE_NOT_GPU", + "gpu_count": world, + "gpu_snapshot_sha256": "a" * 64, + "max_draft_len": k, + "rank_local_graph_batch_sizes": list(graphs), + "runtime_snapshot": "CPU_FIXTURE_NOT_NATIVE", + "source_diff_sha256": "b" * 64, + "source_head": "ORIGINAL_RECORDED_CPU_FIXTURE_SOURCE", + "topology": "CPU_FIXTURE_NOT_FABRIC", + } + tables, measurements = {}, [] + for g in graphs: + budgets = [0] if k == 1 else [0, g, g * 2] + values = [float(index + 1) for index in range(len(budgets))] + tables[str(g)] = {"token_counts": budgets, "step_time_ms": values} + measurements.extend( + { + "rank_local_graph_batch_size": g, + "rank_local_verifier_budget": v, + "source_result_sha256": "c" * 64, + "step_time_ms": value, + } + for v, value in zip(budgets, values) + ) + return { + "schema_version": 2, + "cost_tables": tables, + "engine_fingerprint": fingerprint, + "engine_fingerprint_sha256": _digest(fingerprint), + "minimum_predicted_gain": 0.01, + "measurements": measurements, + } + + +def _check(payload: dict, k: int = 3, graphs: tuple[int, ...] = (16, 64)) -> dict: + return MODULE.check_cost_bytes( + json.dumps(payload).encode(), physical_k=k, graph_batch_sizes=list(graphs) + ) + + +class PublicCostExportTests(unittest.TestCase): + def _whole_step_summary(self, k: int = 3) -> dict: + """Exercise the real reducer on explicitly synthetic CPU raw records.""" + utilities = MODULE._source_module( + "cost_bridge_test_reducer", ROOT / "tensorrt_llm/tools/verify_costs.py" + ) + costs = _costs(k=k, graphs=(16,), world=2) + declared = [] + shapes = [] + for v in costs["cost_tables"]["16"]["token_counts"]: + shape = { + "workload": "dspark_chain", + "physical_k": k, + "graph_batch_size": 16, + "verifier_budget": v, + "attention_layout": "uniform" if v == 0 else "packed_chain_rows", + "captured_head": True, + } + shapes.append(shape) + declared.append(utilities.canonical_sha256(shape)) + identity = { + "cpu_fixture_not_gpu": True, + "generations_per_shape": 3, + "declared_shape_ids": declared, + "recorded_engine_fingerprint": costs["engine_fingerprint"], + "profile": { + "physical_k": k, + "record_confidence_head": True, + "measurement_scope": "whole-step", + "cells": [ + {"graph_batch_size": 16, "verifier_budget": shape["verifier_budget"]} + for shape in shapes + ], + }, + } + raw = [] + for cell_index, (shape, declared_id) in enumerate(zip(shapes, declared)): + for generation in range(3): + for rank in range(2): + raw.append( + { + "schema": "verify-cost-raw-v1", + "rank": rank, + "world_size": 2, + "generation": generation, + "generation_closed": True, + "run_complete": True, + "model_shutdown_completed": True, + "identity": identity, + "timing_scope": MODULE.WHOLE_STEP_SCOPE, + "declared_shape_id": declared_id, + "execution_shape": {**shape, "observed_graph_key": [16, k]}, + "samples": [ + { + "sample_index": cell_index * 42 + generation * 14 + i, + "executor_iteration": cell_index * 42 + + generation * 14 + + i + + 100, + "elapsed_ms": float(rank + i + cell_index + 1), + "eligible": True, + } + for i in range(14) + ], + } + ) + return utilities.reduce_measurements(raw) + + def test_complete_loop_reducer_to_actual_consumer_preserves_values(self) -> None: + for k in (1, 2, 3, 5, 8): + summary = self._whole_step_summary(k) + before = copy.deepcopy(summary) + data = json.dumps(summary).encode() + output, report = MODULE.runtime_cost_bytes(data, physical_k=k, graph_batch_sizes=[16]) + payload = json.loads(output) + self.assertEqual(summary, before) + self.assertEqual( + payload["engine_fingerprint"], summary["identity"]["recorded_engine_fingerprint"] + ) + self.assertTrue(report["direct_whole_step_measurements_exported"]) + self.assertFalse(report["performance_or_economic_admission_qualified"]) + self.assertNotIn("iteration_drain_guard", payload) + expected = { + cell["execution_shape"]["verifier_budget"]: statistics.median( + cell["retained_rank_max_samples_ms"] + ) + for cell in summary["cells"] + } + for row in payload["measurements"]: + self.assertEqual(row["step_time_ms"], expected[row["rank_local_verifier_budget"]]) + self.assertEqual(row["source_result_sha256"], hashlib.sha256(data).hexdigest()) + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "costs.json" + path.write_bytes(output) + consumer = MODULE._source_module( + "bridge_test_actual_consumer", + ROOT / "tensorrt_llm/_torch/speculative/dspark_planner.py", + ) + table, _ = consumer.load_runtime_sps_cost_table( + path, max_draft_len=k, graph_batch_sizes=[16], validate_freshness=False + ) + self.assertEqual(table.step_time(0, 16), expected[0]) + + def test_complete_loop_export_refuses_missing_or_repriced_cells(self) -> None: + for kind in ("missing", "median", "fingerprint", "K", "layout", "duplicate", "floor"): + summary = self._whole_step_summary() + if kind == "missing": + summary["cells"].pop() + elif kind == "median": + summary["cells"][0]["rank_max_median_ms"] += 1.0 + elif kind == "fingerprint": + del summary["identity"]["recorded_engine_fingerprint"] + elif kind == "K": + summary["cells"][0]["execution_shape"]["physical_k"] = 5 + elif kind == "layout": + summary["cells"][0]["execution_shape"]["attention_layout"] = "other" + elif kind == "duplicate": + summary["cells"].append(copy.deepcopy(summary["cells"][0])) + else: + summary["sample_floors"]["pooled_per_shape"] = 1 + with self.subTest(kind=kind), self.assertRaises((ValueError, KeyError)): + MODULE.runtime_cost_bytes( + json.dumps(summary).encode(), physical_k=3, graph_batch_sizes=[16] + ) + + def test_existing_costs_use_actual_consumer_without_freshness_gate(self) -> None: + for k in (1, 2, 3, 4, 5, 8): + for world in (1, 2, 8): + with self.subTest(k=k, world=world): + payload = _costs(k=k, graphs=(8, 192, 256), world=world) + before = copy.deepcopy(payload) + report = _check(payload, k=k, graphs=(8, 192, 256)) + subset_report = _check(payload, k=k, graphs=(8,)) + self.assertTrue(subset_report["structural_validation_passed"]) + self.assertEqual(payload, before) + self.assertTrue(report["structural_validation_passed"]) + self.assertFalse(report["destination_freshness_checked"]) + self.assertFalse(report["new_measurements_created"]) + self.assertFalse(report["performance_or_economic_admission_qualified"]) + self.assertEqual( + report["recorded_source_head"], + payload["engine_fingerprint"]["source_head"], + ) + + def test_optional_measured_drain_values_retained_without_invention(self) -> None: + payload = _costs() + payload["iteration_drain_guard"] = { + "loss_multiplier": 1.5, + "mean_output_tokens_per_request_iteration": 3.25, + "minimum_group_value_ms": 0.2, + "source_result_sha256": "d" * 64, + "tail_graph_batch_size": 16, + } + before = copy.deepcopy(payload) + _check(payload) + self.assertEqual(payload, before) + self.assertNotIn("iteration_drain_guard", _costs()) + + def test_actual_runtime_candidate_bounds_refuse_before_copy(self) -> None: + payload = _costs(graphs=(16,)) + budgets = [0, 16, 17, 18, 19, 20] + payload["cost_tables"]["16"] = { + "token_counts": budgets, + "step_time_ms": [1.0] * len(budgets), + } + payload["measurements"] = [ + { + "rank_local_graph_batch_size": 16, + "rank_local_verifier_budget": v, + "source_result_sha256": "c" * 64, + "step_time_ms": 1.0, + } + for v in budgets + ] + with self.assertRaisesRegex(ValueError, "compact V cells"): + _check(payload, graphs=(16,)) + + def test_actual_consumer_rejects_invalid_schema_integrity_or_cells(self) -> None: + for kind in ( + "physical_k", + "graph_coverage", + "fingerprint", + "measurement", + "native", + "negative", + "extra", + "bool", + "unknown_scope", + ): + payload = _costs() + k, graphs = 3, (16, 64) + if kind == "physical_k": + k = 5 + elif kind == "graph_coverage": + graphs = (16, 32) + elif kind == "fingerprint": + payload["engine_fingerprint_sha256"] = "f" * 64 + elif kind == "measurement": + payload["measurements"][0]["step_time_ms"] = 99.0 + elif kind == "native": + payload["cost_tables"]["16"]["token_counts"][0] = 1 + elif kind == "negative": + payload["cost_tables"]["16"]["step_time_ms"][0] = -1.0 + elif kind == "extra": + payload["measurements"].append(copy.deepcopy(payload["measurements"][0])) + elif kind == "bool": + payload["cost_tables"]["16"]["token_counts"][1] = True + else: + payload["timing_scope"] = "partial-scope-cannot-be-hidden" + with self.subTest(kind=kind), self.assertRaises((TypeError, ValueError)): + _check(payload, k=k, graphs=graphs) + + def test_partial_summary_refuses_even_if_eligibility_is_toggled(self) -> None: + for schema in ("verify-cost-raw-v1", "verify-cost-measured-summary-v1"): + for eligible in (False, True): + payload = { + "schema": schema, + "timing_scope": "original-graph-replay-entry-including-static-copies", + "identity": {"identity_scope": "DIAGNOSTIC_PARTIAL_LOADED_RUNTIME"}, + "production_sps_table_eligible": eligible, + } + with self.subTest(schema=schema, eligible=eligible): + with self.assertRaisesRegex(ValueError, "whole-step producer"): + _check(payload) + + def test_ambiguous_or_nonfinite_json_refuses_before_export(self) -> None: + for data in (b'{"schema_version":2,"schema_version":2}', b'{"value":NaN}', b"[]"): + with self.subTest(data=data), self.assertRaises((ValueError, TypeError)): + MODULE.check_cost_bytes(data, physical_k=3, graph_batch_sizes=[16, 64]) + + def test_real_standalone_cli_preserves_bytes_and_refuses_overwrite(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source, output = root / "costs.json", root / "new-costs.json" + original = json.dumps(_costs(), indent=3).encode() + b"\n\n" + source.write_bytes(original) + command = [ + sys.executable, + "-I", + "-B", + str(SOURCE), + "--input", + str(source), + "--physical-k", + "3", + "--graph-batch-sizes", + "16", + "64", + "--output", + str(output), + ] + result = subprocess.run(command, cwd=root, capture_output=True, timeout=30) + self.assertEqual(result.returncode, 0, result.stderr.decode()) + self.assertEqual(source.read_bytes(), original) + self.assertEqual(output.read_bytes(), original) + report = json.loads(result.stdout) + self.assertEqual(report["source_artifact_sha256"], hashlib.sha256(original).hexdigest()) + again = subprocess.run(command, cwd=root, capture_output=True, timeout=30) + self.assertNotEqual(again.returncode, 0) + self.assertEqual(output.read_bytes(), original) + + def test_real_complete_loop_cli_writes_only_original_direct_cells(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source, output = root / "CPU-summary.json", root / "CPU-costs.json" + summary = self._whole_step_summary(k=5) + source.write_text(json.dumps(summary), encoding="utf-8") + result = subprocess.run( + [ + sys.executable, + "-I", + "-B", + str(SOURCE), + "--input", + str(source), + "--physical-k", + "5", + "--graph-batch-sizes", + "16", + "--output", + str(output), + ], + cwd=root, + capture_output=True, + timeout=30, + ) + self.assertEqual(result.returncode, 0, result.stderr.decode()) + payload = json.loads(output.read_bytes()) + self.assertEqual( + payload["engine_fingerprint"], summary["identity"]["recorded_engine_fingerprint"] + ) + self.assertEqual(payload["minimum_predicted_gain"], 0.01) + self.assertFalse( + json.loads(result.stdout)["performance_or_economic_admission_qualified"] + ) + + def test_real_partial_export_cli_creates_no_cost_artifact(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source, output = root / "partial.json", root / "must-not-exist.json" + source.write_text(json.dumps({"schema": "verify-cost-measured-summary-v1"})) + result = subprocess.run( + [ + sys.executable, + "-I", + "-B", + str(SOURCE), + "--input", + str(source), + "--physical-k", + "3", + "--graph-batch-sizes", + "16", + "64", + "--output", + str(output), + ], + cwd=root, + capture_output=True, + timeout=30, + ) + self.assertNotEqual(result.returncode, 0) + self.assertIn(b"whole-step producer", result.stderr) + self.assertFalse(output.exists()) + + def test_cpu_dependency_loading_never_initializes_native_package(self) -> None: + _check(_costs()) + self.assertNotIn("tensorrt_llm", sys.modules) + self.assertNotIn("torch", sys.modules) + source = SOURCE.read_text(encoding="utf-8") + for token in ("/scratch/", "aws-pdx", "K3", "admitted=True"): + self.assertNotIn(token, source) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/unittest/tools/test_dspark_observations.py b/tests/unittest/tools/test_dspark_observations.py new file mode 100644 index 000000000000..9fbc2367f47b --- /dev/null +++ b/tests/unittest/tools/test_dspark_observations.py @@ -0,0 +1,445 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Source/CPU controls only; finite parents do not qualify native overlap/CUDA.""" + +import copy +import importlib.util +import json +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace + +import torch + +ROOT = Path(__file__).resolve().parents[3] +PUBLIC = ROOT / "tensorrt_llm/tools" + + +def load(name, path, source=None): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + if source is None: + spec.loader.exec_module(module) + else: + exec(compile(source, str(path), "exec"), module.__dict__) + return module + + +observations = load("tensorrt_llm.tools.dspark_observations", PUBLIC / "dspark_observations.py") +RUNTIME = ROOT / "tensorrt_llm/_torch/pyexecutor/dspark_observation_runtime.py" +runtime = load("_public_observation_runtime", RUNTIME) + + +def request(identifier, slot): + return SimpleNamespace( + py_request_id=identifier, + py_seq_slot=slot, + py_decoding_iter=0, + is_dummy=False, + py_draft_tokens=[7, 8, 9], + ) + + +class FiniteSampler: + """Deferred copies covered by the finite ORIGINAL sampler event, not a GPU event.""" + + def __init__(self): + self.copies = [] + self.waits = 0 + + def _copy_to_host(self, tensor): + host = torch.zeros_like(tensor, device="cpu") + self.copies.append((host, tensor.clone())) + return host + + def event(self): + copies, self.copies = self.copies, [] + + def synchronize(): + self.waits += 1 + for host, snapshot in copies: + host.copy_(snapshot) + + return SimpleNamespace(synchronize=synchronize) + + +class FiniteWorker: + def __init__(self): + self.sequence = 0 + self.confidence = torch.zeros((4, 3)) + self.stamps = torch.zeros(4, dtype=torch.int32) + + def bump_draft_seq(self): + old = self.sequence + self.sequence += 1 + return old + + def staged_confidence_buffer(self): + return self.confidence + + def confidence_stamp_buffer(self): + return self.stamps + + def confidence_row_for(self, identifier): + return {41: 0, 42: 3}[identifier] + + +class FiniteOriginal: + """Only the parent's finite call/order/buffer contract is modeled here.""" + + def __init__(self): + self.is_warmup = False + self.iter_counter = 10 + self.worker = FiniteWorker() + self.model_engine = SimpleNamespace( + _collect_calibration=True, + model=SimpleNamespace(spec_worker=self.worker), + spec_config=SimpleNamespace(max_draft_len=3), + ) + self.sampler = FiniteSampler() + self.tokens = torch.zeros((3, 3), dtype=torch.int32) + + def _forward_step(self, batch, new_tensors_device=None): + for req in batch.generation_requests: + row = self.worker.confidence_row_for(req.py_request_id) + self.worker.confidence[row] = torch.tensor( + [ + self.iter_counter + row, + self.iter_counter + row + 0.1, + self.iter_counter + row + 0.2, + ] + ) + self.worker.stamps[row] = self.worker.sequence + self.tokens[req.py_seq_slot] = torch.tensor( + [self.iter_counter * 10 + i for i in range(3)] + ) + return {"next_draft_tokens": self.tokens} + + def _sample_async(self, batch, output): + host = SimpleNamespace( + next_draft_tokens=self.sampler._copy_to_host(self.tokens), + new_tokens_lens=self.sampler._copy_to_host(torch.tensor([2, 1, 3])), + verify_lens=None, + ) + return SimpleNamespace( + host=host, + device=SimpleNamespace(next_draft_tokens=self.tokens), + sampler_event=self.sampler.event(), + requests=list(batch.generation_requests), + runtime_draft_len=3, + draft_lens=[3] * len(batch.generation_requests), + ) + + def _update_requests(self, state, resource_manager=None): + state.sampler_event.synchronize() + for req in state.requests: + if getattr(req, "retired", False): + continue + req.py_decoding_iter += 1 + req.py_draft_tokens = state.host.next_draft_tokens[req.py_seq_slot].tolist() + + +def complete_overlap(mixin): + class Executor(mixin, FiniteOriginal): + pass + + executor = Executor() + reqs = [request(41, 2), request(42, 0)] + batch = SimpleNamespace(generation_requests=reqs) + output = executor._forward_step(batch) + source = executor._sample_async(batch, output) + executor.iter_counter = 11 + reordered = SimpleNamespace(generation_requests=list(reversed(reqs))) + output = executor._forward_step(reordered, source.device) + executor._update_requests(source) + target = executor._sample_async(reordered, output) + executor._update_requests(target) + return executor + + +def run_record(mode): + inputs = [{"input_id": 0, "prompt_token_ids": [4, 5]}] + return { + "schema": "dspark-public-run-v1", + "mode": mode, + "run_complete": True, + "model_shutdown_completed": True, + "calibration_collection": False, + "comparison_contract": { + "physical_k": 3, + "batch_size": 1, + "output_tokens": 2, + "same_cpu_fixture": True, + }, + "elapsed_seconds": 2 if mode == "static" else 1.5, + "resolved_confidence_enabled": mode == "dynamic", + "runtime_confidence_enabled": mode == "dynamic", + "routing": {"native_replays": 2, "compact_replays": 1 if mode == "dynamic" else 0}, + "inputs": inputs, + "responses": [{**inputs[0], "output_token_ids": [6, 7], "finish_reason": "length"}], + } + + +class Tests(unittest.TestCase): + def test_whole_overlap_reordered_slots_and_delayed_original_event(self): + executor = complete_overlap(runtime.ObservationExecutorMixin) + payload = executor._native_pairs.payload() + self.assertEqual(executor.sampler.waits, 2) + self.assertEqual( + payload["logits"].tolist()[0], [13.0, 13.100000381469727, 13.199999809265137] + ) + self.assertEqual( + payload["logits"].tolist()[1], [10.0, 10.100000381469727, 10.199999809265137] + ) + self.assertEqual(payload["prefix_mask"].tolist(), [[1, 0, 0], [1, 1, 0]]) + self.assertTrue(all(join["pair_key"][2:] == [10, 11] for join in payload["meta"]["joins"])) + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "pairs.pt" + torch.save(payload, path) + loaded = torch.load(path, map_location="cpu", weights_only=True) + self.assertEqual(loaded["meta"]["pairing"], "draft_seq_ring") + + def test_whole_old_fixed_restored_persistent_input_overwrite(self): + source = RUNTIME.read_text(encoding="utf-8") + before = "input_host = self.sampler._copy_to_host(new_tensors_device.next_draft_tokens)" + old = source.replace(before, "input_host = new_tensors_device.next_draft_tokens", 1) + for index, variant in enumerate((old, source, old)): + loaded = load(f"_source_boundary_{index}", RUNTIME, variant) + if index == 1: + self.assertEqual( + complete_overlap(loaded.ObservationExecutorMixin)._native_pairs.stats[ + "recorded" + ], + 2, + ) + else: + with self.assertRaisesRegex(ValueError, "verification input differs"): + complete_overlap(loaded.ObservationExecutorMixin) + + def test_wrong_step_incarnation_stale_confidence_and_censored_negatives(self): + executor = complete_overlap(runtime.ObservationExecutorMixin) + source = next(iter(executor._observation_states.values())) + req = source.requests[41][0] + arguments = dict( + source=source, + request=req, + target_iteration=12, + input_proposal=source.state.host.next_draft_tokens[2], + accepted=2, + verified=3, + target_slot=2, + ) + for fields, message in ( + ({"target_iteration": 10}, "after"), + ({"request": request(41, 2)}, "incarnation"), + ({"verified": 2}, "censors"), + ({"accepted": 4}, "outside"), + ): + with self.assertRaisesRegex(ValueError, message): + observations.NativePairs(3).record(**{**arguments, **fields}) + source.stamp_host[0] += 1 + with self.assertRaisesRegex(ValueError, "producing draft"): + observations.NativePairs(3).record(**arguments) + + def test_retired_target_is_not_labeled(self): + class Executor(runtime.ObservationExecutorMixin, FiniteOriginal): + pass + + executor = Executor() + req = request(41, 2) + batch = SimpleNamespace(generation_requests=[req]) + state = executor._sample_async(batch, executor._forward_step(batch)) + executor._update_requests(state) + executor.iter_counter += 1 + output = executor._forward_step(batch, state.device) + target = executor._sample_async(batch, output) + req.retired = True + executor._update_requests(target) + self.assertEqual(executor._native_pairs.stats["recorded"], 0) + self.assertEqual(executor._native_pairs.stats["retired_or_not_advanced"], 1) + + def test_same_state_cannot_update_twice(self): + executor = complete_overlap(runtime.ObservationExecutorMixin) + state = list(executor._observation_states.values())[-1].state + with self.assertRaisesRegex(RuntimeError, "consumed"): + executor._update_requests(state) + + def test_configurable_physical_k(self): + for k in (1, 2, 3, 4, 5, 8): + req = request(41, 0) + tokens = torch.arange(k).view(1, -1) + state = SimpleNamespace(host=SimpleNamespace(next_draft_tokens=tokens)) + snapshot = observations.ProposalSnapshot( + 1, state, torch.zeros((1, k)), {41: (req, 0, 0)}, torch.tensor([8]), 8, True + ) + pairs = observations.NativePairs(k) + pairs.record( + source=snapshot, + request=req, + target_iteration=2, + input_proposal=tokens[0], + accepted=k - 1, + verified=k, + target_slot=0, + ) + self.assertEqual(pairs.payload()["prefix_mask"].tolist(), [[1] * (k - 1) + [0]]) + + def test_prefill_placeholder_is_not_labeled_as_scored_draft(self): + class Executor(runtime.ObservationExecutorMixin, FiniteOriginal): + pass + + executor = Executor() + req = request(41, 2) + empty = SimpleNamespace(generation_requests=[]) + source = executor._sample_async(empty, executor._forward_step(empty)) + executor.iter_counter += 1 + batch = SimpleNamespace(generation_requests=[req]) + output = executor._forward_step(batch, source.device) + executor._update_requests(source) + target = executor._sample_async(batch, output) + executor._update_requests(target) + self.assertEqual(executor._native_pairs.stats["recorded"], 0) + self.assertEqual(executor._native_pairs.stats["prefill_placeholder_no_scored_proposal"], 1) + + def test_literal_collect_payload_to_approved_public_fitter(self): + fitter = PUBLIC / "dspark_calibration.py" + self.assertTrue(fitter.is_file(), "Assemble the approved public fitter with this endpoint") + with tempfile.TemporaryDirectory() as directory: + cwd = Path(directory) + shard = cwd / "cpu-source-contract-pairs.pt" + torch.save( + complete_overlap(runtime.ObservationExecutorMixin)._native_pairs.payload(), shard + ) + output = cwd / "cpu-source-contract-temperatures.json" + result = subprocess.run( + [ + sys.executable, + "-I", + "-B", + str(fitter), + "--shards", + str(shard), + "--physical-k", + "3", + "--minimum-samples", + "2", + "--output", + str(output), + ], + cwd=cwd, + capture_output=True, + text=True, + timeout=60, + ) + self.assertEqual(result.returncode, 0, result.stderr) + payload = json.loads(output.read_bytes()) + self.assertEqual(payload["physical_k"], 3) + self.assertEqual(payload["num_samples"], 2) + self.assertFalse(payload["task_accuracy_or_economic_admission_qualified"]) + + def test_complete_public_run_parser_has_configurable_k_bs_and_no_internal_kwargs(self): + cli = load("_public_cli_parser_control", PUBLIC / "dspark_reproduction.py") + name = "tensorrt_llm._torch.speculative.dspark_reproduction" + original = sys.modules.get(name) + seen = [] + sys.modules[name] = SimpleNamespace(run_native=lambda args: seen.append(args)) + try: + for mode in ("static", "dynamic", "collect"): + cli.main( + [ + "run", + "--mode", + mode, + "--model", + "model", + "--config", + "config.yml", + "--inputs", + "input.jsonl", + "--source-root", + "checkout", + "--physical-k", + "4", + "--batch-size", + "128", + "--output-tokens", + "64", + "--timeout", + "25200", + "--output", + "private-output", + ] + ) + self.assertEqual([args.mode for args in seen], ["static", "dynamic", "collect"]) + self.assertTrue(all(args.physical_k == 4 and args.batch_size == 128 for args in seen)) + self.assertTrue(all("executor_cls" not in vars(args) for args in seen)) + finally: + if original is None: + del sys.modules[name] + else: + sys.modules[name] = original + + def test_complete_literal_cli_compare_and_no_clobber(self): + with tempfile.TemporaryDirectory() as directory: + cwd = Path(directory) + files = [] + for mode in ("static", "dynamic"): + path = cwd / f"{mode}.json" + path.write_text(json.dumps(run_record(mode)), encoding="utf-8") + files.append(path) + output = cwd / "comparison.json" + command = [ + sys.executable, + "-I", + "-B", + str(PUBLIC / "dspark_reproduction.py"), + "compare", + "--runs", + *map(str, files), + "--output", + str(output), + ] + result = subprocess.run(command, cwd=cwd, capture_output=True, text=True, timeout=60) + self.assertEqual(result.returncode, 0, result.stderr) + report = json.loads(output.read_bytes()) + self.assertAlmostEqual(report["dynamic_over_static_pct"], 100 / 3) + self.assertIsNone(report["confidence_interval"]) + self.assertFalse(report["task_accuracy_qualified"]) + self.assertNotEqual( + subprocess.run(command, cwd=cwd, capture_output=True, timeout=60).returncode, 0 + ) + + def test_static_fallback_no_compact_unmatched_and_partial_refused(self): + static, dynamic = run_record("static"), run_record("dynamic") + for field, value, message in ( + ("resolved_confidence_enabled", False, "fell back"), + ("runtime_confidence_enabled", False, "fell back"), + ("routing", {"native_replays": 3, "compact_replays": 0}, "never executed"), + ("run_complete", False, "closed"), + ): + bad = {**dynamic, field: value} + with self.assertRaisesRegex(ValueError, message): + observations.compare_runs([static, bad]) + different = copy.deepcopy(dynamic) + different["comparison_contract"]["physical_k"] = 4 + with self.assertRaisesRegex(ValueError, "differ"): + observations.compare_runs([static, different]) + different = copy.deepcopy(dynamic) + different["responses"][0]["prompt_token_ids"] = [999] + with self.assertRaisesRegex(ValueError, "prompt"): + observations.compare_runs([static, different]) + + def test_internal_constructor_kwargs_not_public_llm_args(self): + # Real argument-validation and constructor controls live in the separate + # native-runtime test. This CPU guard prevents accidental private API injection. + driver = (ROOT / "tensorrt_llm/_torch/speculative/dspark_reproduction.py").read_text() + self.assertNotIn("executor_cls", driver) + self.assertNotIn("runtime_constructors", driver) + + +if __name__ == "__main__": + unittest.main()