From 52ab4df64b9397f92683ab179ea6e2f6c7b3d3a0 Mon Sep 17 00:00:00 2001 From: Zhaoyuan He Date: Tue, 29 Sep 2026 14:51:44 -0700 Subject: [PATCH 1/3] [TRTLLM-14620][feat] Add ragged DSpark verification semantics Represent per-request verification windows consistently in packed layouts, sampling metadata, occurrence penalties, and completed-step rewind accounting. Match synthetic proposals in target vocabulary to preserve request-local bonus distributions without mutating stored probabilities. Keep uniform dispatch when ragged metadata is absent and leave attention kernels, confidence scheduling, and model-engine enablement outside this change. Signed-off-by: Zhaoyuan He --- tensorrt_llm/_torch/pyexecutor/llm_request.py | 19 + .../_torch/pyexecutor/sampler/penalties.py | 76 +- .../_torch/speculative/dspark_ragged.py | 463 ++++++++++++ tensorrt_llm/_torch/speculative/interface.py | 239 +++++- .../_torch/speculative/spec_sampler_base.py | 102 ++- .../test_kv_cache_v2_capacity_only.py | 39 +- .../hw_agnostic/test_dspark_ragged.py | 500 ++++++++++++ .../test_dspark_ragged_metadata.py | 90 +++ .../hw_agnostic/test_ragged_spec_sampler.py | 714 ++++++++++++++++++ .../test_capture_sampling_params.py | 6 +- .../speculative/test_group_all_greedy_sync.py | 1 + .../test_rejection_buffers_guard.py | 2 + 12 files changed, 2196 insertions(+), 55 deletions(-) create mode 100644 tensorrt_llm/_torch/speculative/dspark_ragged.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged_metadata.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index f3edf2e3cbff..1eb71268ae98 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -1717,3 +1717,22 @@ def get_draft_token_length(request: LlmRequest) -> int: if request.py_draft_tokens is not None: return len(request.py_draft_tokens) return 0 + + +def get_request_tokens_per_gen_step(request: LlmRequest, + runtime_tokens_per_gen_step: int) -> int: + """How many tokens this request contributes to one generation step. + + Uniform speculation gives every request the batch-wide + ``runtime_tokens_per_gen_step`` (1 accepted token + ``runtime_draft_len`` + draft positions). Ragged verification -- where a scheduler hands each + request its own verify window -- records the per-request draft length in + ``py_verify_len``. + + ``py_verify_len`` is set only by the ragged path, so every other + speculation mode gets the batch-wide value back unchanged. + """ + verify_len = getattr(request, "py_verify_len", None) + if verify_len is None: + return runtime_tokens_per_gen_step + return 1 + int(verify_len) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/penalties.py b/tensorrt_llm/_torch/pyexecutor/sampler/penalties.py index 7edfc6818a96..bc498e486b22 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/penalties.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/penalties.py @@ -672,6 +672,12 @@ class SpecMetadataLike(Protocol): vocab_size: int penalty_state: Optional["PenaltyState"] batch_slot_ids: Optional[torch.Tensor] + verify_lens: Optional[torch.Tensor] + qo_indptr: Optional[torch.Tensor] + total_verify_tokens: Optional[int] + + @property + def is_ragged_verify(self) -> bool: ... @dataclass(kw_only=True) @@ -886,13 +892,16 @@ def build_row_mapping( draft_len: int, draft_tokens: torch.Tensor, device: torch.device, + num_logit_rows: Optional[int] = None, ) -> Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: """Derive the per-row slot map and intra-step prefix for a packed logits batch. - Assumes the layout every linear one-model mode produces once - ``_reshape_logits_for_accept`` has run: ``[ctx (1 row each), gen (draft_len + 1 - rows each)]``, request-major. PARD emits ``2 * draft_len`` rows per gen request - but that hook already narrows them to ``draft_len + 1``, so both share this code. + Uniform linear modes produce ``[ctx (1 row each), gen (draft_len + 1 rows + each)]`` once ``_reshape_logits_for_accept`` has run. Ragged verification + instead packs each generation request's ``verify_lens[r]`` token rows back to + back on the flat token axis. In either layout, rows for one request are + contiguous and ordered by speculative position. PARD emits ``2 * draft_len`` + rows per gen request but its reshape hook narrows them before this function. Tree modes do NOT fit: their rows are tree nodes, so a row's prefix is its root path rather than the rows before it. They are rejected at admission. @@ -909,16 +918,56 @@ def build_row_mapping( return None num_gens = batch_size - num_contexts - rows_per_gen = draft_len + 1 - total_rows = num_contexts + num_gens * rows_per_gen - - batch_slots = slot_ids[:batch_size].to(torch.int64) - # Context requests own one row; gen requests own rows_per_gen consecutive rows. + batch_slots = slot_ids[:batch_size].to(device=device, dtype=torch.int64) + + if getattr(spec_metadata, "is_ragged_verify", False): + from tensorrt_llm._torch.speculative.dspark_ragged import row_ids_from_lens + + verify_lens = getattr(spec_metadata, "verify_lens", None) + qo_indptr = getattr(spec_metadata, "qo_indptr", None) + total_verify_tokens = getattr(spec_metadata, "total_verify_tokens", None) + if ( + verify_lens is None + or qo_indptr is None + or verify_lens.dim() != 1 + or qo_indptr.dim() != 1 + or verify_lens.numel() != num_gens + or qo_indptr.numel() != num_gens + 1 + ): + return None + if total_verify_tokens is None: + if num_logit_rows is None: + return None + total_verify_tokens = num_logit_rows - num_contexts + total_gen_rows = int(total_verify_tokens) + if total_gen_rows < 0: + return None + total_rows = num_contexts + total_gen_rows + if num_logit_rows is not None and total_rows != num_logit_rows: + return None + + lens = verify_lens.to(device=device, dtype=torch.int64) + gen_request_ids = row_ids_from_lens(lens, total=total_gen_rows) + packed_positions = torch.arange(total_gen_rows, device=device) + starts = qo_indptr[:num_gens].to(device=device, dtype=torch.int64) + position_in_request = packed_positions - starts[gen_request_ids] + else: + rows_per_gen = draft_len + 1 + total_gen_rows = num_gens * rows_per_gen + total_rows = num_contexts + total_gen_rows + if num_logit_rows is not None and total_rows != num_logit_rows: + return None + gen_rows = torch.arange(total_gen_rows, device=device) + gen_request_ids = gen_rows // rows_per_gen + position_in_request = gen_rows % rows_per_gen + + # Context requests own one row. Generation rows are request-contiguous in + # both the uniform rectangle and the ragged packed token axis. row_slots = ( torch.cat( [ batch_slots[:num_contexts], - batch_slots[num_contexts:batch_size].repeat_interleave(rows_per_gen), + batch_slots[num_contexts:batch_size][gen_request_ids], ] ) if num_gens > 0 @@ -935,15 +984,12 @@ def build_row_mapping( intra_valid = torch.zeros((total_rows, max(draft_len, 0)), dtype=torch.bool, device=device) if num_gens > 0 and draft_len > 0 and draft_tokens.numel() > 0: drafts = draft_tokens.reshape(num_gens, -1)[:, :draft_len].to(torch.int64) - gen_rows = torch.arange(num_gens * rows_per_gen, device=device) - g_of_row = gen_rows // rows_per_gen # which request owns the row - k_of_row = gen_rows % rows_per_gen # the row's speculative position # Every gen row carries its own request's drafts; validity alone selects the # strictly-earlier prefix. Column j counts for row k iff j < k. - intra_tokens[num_contexts:] = drafts[g_of_row] + intra_tokens[num_contexts:] = drafts[gen_request_ids] intra_valid[num_contexts:] = torch.arange(draft_len, device=device).unsqueeze( 0 - ) < k_of_row.unsqueeze(1) + ) < position_in_request.unsqueeze(1) return row_slots, intra_tokens, intra_valid diff --git a/tensorrt_llm/_torch/speculative/dspark_ragged.py b/tensorrt_llm/_torch/speculative/dspark_ragged.py new file mode 100644 index 000000000000..3c8b6d7826ec --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_ragged.py @@ -0,0 +1,463 @@ +# 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. +"""Ragged verification layout: per-request verify lengths in one packed batch, +as a flat token axis plus a ``qo_indptr`` exclusive prefix sum. + +Two invariants keep the pipeline safe: ``verify_lens >= 1`` for every request +(position 0 carries the bonus/anchor token; a request that verifies nothing +never progresses), and ``sum(verify_lens) == graph_num_tokens`` after +:meth:`RaggedVerifyLayout.fill_bucket` -- ``seq_lens.sum()`` becomes +``attn_metadata.num_tokens``, which is all-gathered across attention-DP ranks +and drives the MoE's chunk count, and an under-filled bucket desynchronizes +them without raising. +""" + +from collections.abc import Sequence +from dataclasses import dataclass + +import torch + +__all__ = [ + "RaggedVerifyLayout", + "build_qo_indptr", + "row_ids_from_lens", + "scatter_ragged_to_padded", + "count_accepted_ragged", + "fill_padded_rows_onehot", + "ragged_gather_index_lists", + "RaggedPadSplit", + "resolve_ragged_pad_split", + "fill_bucket_device", + "build_row_maps_device", +] + + +def build_qo_indptr(verify_lens: torch.Tensor) -> torch.Tensor: + """``[bs]`` lengths -> ``[bs + 1]`` exclusive prefix sum, on the same device. + + Kept as a tensor op with no host sync so it can run inside a captured graph. + """ + if verify_lens.dim() != 1: + raise ValueError(f"verify_lens must be 1-D, got {tuple(verify_lens.shape)}") + lens = verify_lens.to(torch.int32) + out = torch.zeros(lens.numel() + 1, dtype=torch.int32, device=lens.device) + torch.cumsum(lens, dim=0, out=out[1:]) + return out + + +@dataclass +class RaggedVerifyLayout: + """One packed ragged verify batch. + + Attributes: + verify_lens: ``[bs]`` int32, each >= ``min_verify_len``. + qo_indptr: ``[bs + 1]`` int32 exclusive prefix sum of ``verify_lens``. + extend_start_loc: ``[bs]`` int32, ``qo_indptr[:-1]``. + graph_num_tokens: the captured token bucket this batch is padded to. + total_verify_tokens: ``sum(verify_lens)``, host-side when known. + num_real_requests: how many leading rows are real requests. Rows beyond + it are padding added to reach the captured batch size; their outputs + must be discarded. Equals ``bs`` until :meth:`fill_bucket` pads. + """ + + verify_lens: torch.Tensor + qo_indptr: torch.Tensor + extend_start_loc: torch.Tensor + graph_num_tokens: int + total_verify_tokens: int | None = None + num_real_requests: int | None = None + + @property + def bs(self) -> int: + return int(self.verify_lens.shape[0]) + + @property + def num_pad_requests(self) -> int: + return 0 if self.num_real_requests is None else self.bs - self.num_real_requests + + @classmethod + def from_verify_lens( + cls, + verify_lens: torch.Tensor, + *, + graph_num_tokens: int, + total_verify_tokens: int | None = None, + ) -> "RaggedVerifyLayout": + """Pack ``[bs]`` lengths against an already selected graph token count.""" + lens = verify_lens.to(torch.int32) + indptr = build_qo_indptr(lens) + return cls( + verify_lens=lens, + qo_indptr=indptr, + extend_start_loc=indptr[:-1], + graph_num_tokens=int(graph_num_tokens), + total_verify_tokens=( + int(total_verify_tokens) if total_verify_tokens is not None else None + ), + ) + + def fill_bucket( + self, *, max_verify_len: int, padded_bs: int | None = None + ) -> "RaggedVerifyLayout": + """Pad rows to ``padded_bs`` and tokens to exactly ``graph_num_tokens`` + (a short batch desynchronizes attention from the MoE; see the module + docstring). Slack goes to real requests first (up to ``max_verify_len``), + then to pad rows; every pad row needs at least one token or its empty + ``qo_indptr`` range breaks per-row slicing. Raises when the bucket + cannot be hit exactly. + """ + if self.total_verify_tokens is None: + raise ValueError("fill_bucket needs the host-side total_verify_tokens") + real = [int(v) for v in self.verify_lens.tolist()] + n_real = len(real) + padded_bs = n_real if padded_bs is None else int(padded_bs) + if padded_bs < n_real: + raise ValueError( + f"padded_bs {padded_bs} < real request count {n_real}: the " + f"captured batch size cannot hold this batch" + ) + n_pad = padded_bs - n_real + + # Every pad row must carry at least one token. + baseline = sum(real) + n_pad + capacity = padded_bs * max_verify_len + if baseline > self.graph_num_tokens: + raise ValueError( + f"bucket {self.graph_num_tokens} is too small: {sum(real)} real " + f"tokens plus {n_pad} pad rows already needs {baseline}" + ) + if self.graph_num_tokens > capacity: + raise ValueError( + f"bucket {self.graph_num_tokens} exceeds what {padded_bs} rows can " + f"absorb at max_verify_len={max_verify_len} ({capacity}); pick a " + f"denser bucket grid or a larger padded_bs" + ) + + lens = real + [1] * n_pad + spare = self.graph_num_tokens - baseline + # Round-robin: spread extra tokens across real rows (not one request's + # low-survival tail), then spill to pad rows. + for lo, hi in ((0, n_real), (n_real, padded_bs)): + i = lo + while spare > 0 and any(lens[j] < max_verify_len for j in range(lo, hi)): + if lens[i] < max_verify_len: + lens[i] += 1 + spare -= 1 + i = lo + (i - lo + 1) % max(hi - lo, 1) + assert spare == 0, f"internal: {spare} tokens left after filling {lens}" + + filled = torch.tensor(lens, dtype=torch.int32, device=self.verify_lens.device) + out = RaggedVerifyLayout.from_verify_lens( + filled, + graph_num_tokens=self.graph_num_tokens, + total_verify_tokens=sum(lens), + ) + out.num_real_requests = n_real + return out + + def validate( + self, + *, + min_verify_len: int = 1, + max_verify_len: int | None = None, + exact_fill: bool = False, + ) -> None: + """Assert the layout invariants (host-syncing: tests/debug only, never + the hot path). + + ``exact_fill`` additionally requires the token count to hit the bucket + exactly -- the post-:meth:`fill_bucket` contract for anything heading + into a captured graph; a short batch raises nowhere downstream. + """ + lens = [int(v) for v in self.verify_lens.tolist()] + if not lens: + raise ValueError("a ragged layout needs at least one request") + if min(lens) < max(min_verify_len, 1): + raise ValueError( + f"every request must verify its anchor token " + f"(verify_len >= {max(min_verify_len, 1)}), got {lens}" + ) + if max_verify_len is not None and max(lens) > max_verify_len: + raise ValueError( + f"verify_len exceeds the drafted block: {max(lens)} > {max_verify_len}" + ) + total = sum(lens) + if self.total_verify_tokens is not None and total != self.total_verify_tokens: + raise ValueError(f"total_verify_tokens {self.total_verify_tokens} != sum {total}") + if total > self.graph_num_tokens: + raise ValueError( + f"packed tokens {total} exceed the captured bucket {self.graph_num_tokens}" + ) + if exact_fill and total != self.graph_num_tokens: + raise ValueError( + f"packed tokens {total} do not fill the captured bucket " + f"{self.graph_num_tokens}; seq_lens.sum() would disagree with " + f"the row count the graph was captured for" + ) + expected = build_qo_indptr(self.verify_lens) + if not torch.equal(self.qo_indptr, expected): + raise ValueError("qo_indptr is not the prefix sum of verify_lens") + + +@dataclass(frozen=True) +class RaggedPadSplit: + """A shared pad-row window and the token target left for real rows.""" + + pad_len: int + real_target: int + + +def resolve_ragged_pad_split( + *, + bucket: int, + num_real_requests: int, + total_real_tokens: int, + padded_bs: int, + max_verify_len: int, + fixed_pad_len: int | None = None, +) -> RaggedPadSplit | None: + """Resolve one feasible decomposition of a captured ragged bucket. + + All pad rows share one request object and therefore one verify length. The + returned split keeps that length in ``[1, max_verify_len]`` and leaves the + real rows a target between their current token floor and capacity. + """ + num_pad_requests = int(padded_bs) - int(num_real_requests) + if num_pad_requests < 0: + raise ValueError("padded_bs cannot be smaller than num_real_requests") + if max_verify_len < 1: + raise ValueError("max_verify_len must be at least one") + + real_capacity = int(num_real_requests) * int(max_verify_len) + real_floor = int(total_real_tokens) + if num_pad_requests == 0: + pad_len = 1 + elif fixed_pad_len is not None: + pad_len = int(fixed_pad_len) + else: + lower = max(1, -(-(int(bucket) - real_capacity) // num_pad_requests)) + upper = min( + int(max_verify_len), + (int(bucket) - real_floor) // num_pad_requests, + ) + if lower > upper: + return None + pad_len = lower + + if not 1 <= pad_len <= int(max_verify_len): + return None + real_target = int(bucket) - num_pad_requests * pad_len + if not real_floor <= real_target <= real_capacity: + return None + return RaggedPadSplit(pad_len=pad_len, real_target=real_target) + + +def fill_bucket_device( + verify_lens: torch.Tensor, + *, + num_real: torch.Tensor, + graph_num_tokens: int, + max_verify_len: int, + pad_fill: int | None = None, +) -> torch.Tensor: + """Capture-safe :meth:`RaggedVerifyLayout.fill_bucket`: same allocation, + computed entirely on device. + + The host round-robin decomposes into full cycles -- every row with + headroom gains one token -- plus one final partial cycle where the first + ``spare`` headroom rows in index order gain one. Each cycle is therefore + ``grant = headroom & (cumsum(headroom) <= spare)``, and a row's headroom + is at most ``max_verify_len - 1``, so a statically-unrolled loop of that + many cycles per phase (real rows first, then pad rows) reproduces the + host result token-for-token with no ``.item()`` or host sync. + + The device version cannot raise, so feasibility is the caller's contract, + checkable from host-known quantities alone before launch: + ``padded_bs <= graph_num_tokens <= padded_bs * max_verify_len`` and the + scheduled budget small enough that real tokens plus one-per-pad-row fit + the bucket. Under those bounds the result sums to exactly + ``graph_num_tokens``. + + Args: + verify_lens: ``[padded_bs]`` scheduled lengths; entries at or beyond + ``num_real`` are ignored (pad rows restart from one token). + num_real: 0-d integer tensor on the same device (host-known per step, + but not a capture constant -- it changes across replays). + graph_num_tokens: the captured token bucket (capture constant). + max_verify_len: per-row ceiling (capture constant). + pad_fill: when given, every pad row carries EXACTLY this many tokens + and all slack goes to real rows. The host fit publishes this split + (``ragged_pad_verify_len``) because its copy widths -- how many + flat tokens belong to real rows -- must be known before launch; + the device fill has to land on the same split or the host and + device disagree about where real tokens end. + Returns: + ``[padded_bs]`` int32 filled lengths. + """ + if verify_lens.dim() != 1: + raise ValueError(f"verify_lens must be 1-D, got {tuple(verify_lens.shape)}") + padded_bs = verify_lens.numel() + device = verify_lens.device + rows = torch.arange(padded_bs, device=device) + is_real = rows < num_real + lens = torch.where( + is_real, + verify_lens.to(torch.int32), + torch.full( + (padded_bs,), 1 if pad_fill is None else int(pad_fill), dtype=torch.int32, device=device + ), + ) + spare = graph_num_tokens - lens.sum() + phases = (is_real,) if pad_fill is not None else (is_real, ~is_real) + for phase_mask in phases: + for _ in range(max(max_verify_len - 1, 0)): + headroom = phase_mask & (lens < max_verify_len) + in_cycle = torch.cumsum(headroom.to(torch.int64), dim=0) <= spare + grant = (headroom & in_cycle).to(torch.int32) + lens = lens + grant + spare = spare - grant.sum() + return lens + + +def build_row_maps_device( + verify_lens: torch.Tensor, + *, + graph_num_tokens: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Per-token ``(req_idx, kv_correction)`` from filled lengths, on device. + + A request with window ``v`` at within-request offset ``o`` attends KV + extent ``kv_len - v + 1 + o``, i.e. correction ``o - v + 1``. The caller + adds this correction to the owning request's KV length. + + ``verify_lens`` must already sum to ``graph_num_tokens`` + (post-:func:`fill_bucket_device`); the explicit total keeps + ``repeat_interleave`` sync-free. + + Returns: + ``req_idx``: ``[graph_num_tokens]`` int64 owning-row index. + ``correction``: ``[graph_num_tokens]`` int32, in ``[1 - v, 0]``. + """ + indptr = build_qo_indptr(verify_lens).to(torch.long) + req_idx = row_ids_from_lens(verify_lens, total=graph_num_tokens) + offset = torch.arange(graph_num_tokens, device=verify_lens.device) - indptr[req_idx] + correction = (offset - verify_lens.to(torch.long)[req_idx] + 1).to(torch.int32) + return req_idx, correction + + +def ragged_gather_index_lists( + slots: Sequence[int], counts: Sequence[int] +) -> tuple[list[int], list[int]]: + """Row/column index lists for gathering a ragged block from a + ``[num_slots, max_width]`` tensor. + + ``slots[i]`` contributes ``counts[i]`` entries from columns + ``0..counts[i]-1``, concatenated in batch order; with a constant ``counts`` + this matches the uniform ``tensor[slots, :width]`` gather. + """ + if len(slots) != len(counts): + raise ValueError( + f"slots and counts must be the same length, got {len(slots)} and {len(counts)}" + ) + rows: list[int] = [] + cols: list[int] = [] + for slot, count in zip(slots, counts): + if count < 0: + raise ValueError(f"negative gather count {count} for slot {slot}") + rows.extend([slot] * count) + cols.extend(range(count)) + return rows, cols + + +def row_ids_from_lens(verify_lens: torch.Tensor, *, total: int) -> torch.Tensor: + """``[bs]`` lengths -> ``[total]`` owning-request id for each packed token. + + ``total`` (``sum(verify_lens)``) is mandatory: without ``output_size``, + ``repeat_interleave`` with device-resident repeats syncs to size its + output, which is illegal inside the captured graph. + """ + return torch.repeat_interleave( + torch.arange(verify_lens.numel(), device=verify_lens.device), + verify_lens.to(torch.long), + output_size=int(total), + ) + + +def scatter_ragged_to_padded( + flat: torch.Tensor, + *, + verify_lens: torch.Tensor, + qo_indptr: torch.Tensor, + max_len: int, + pad_value: int = 0, +) -> torch.Tensor: + """Unpack a flat ragged batch into ``[bs, max_len, ...]`` with padding. + + Sync-free (capturable): token ``t`` goes to request ``row_ids[t]`` at + column ``t - qo_indptr[row_ids[t]]``, and ``flat.shape[0]`` supplies the + ``output_size`` :func:`row_ids_from_lens` needs. + """ + bs = int(verify_lens.numel()) + rows = row_ids_from_lens(verify_lens, total=flat.shape[0]) + cols = torch.arange(flat.shape[0], device=flat.device) - qo_indptr.to(torch.long)[rows] + out = flat.new_full((bs, max_len, *flat.shape[1:]), pad_value) + out[rows, cols] = flat + return out + + +def fill_padded_rows_onehot( + probs: torch.Tensor, + *, + verify_lens: torch.Tensor, +) -> torch.Tensor: + """Replace ``[bs, max_len, vocab]`` padding rows with a point mass at token 0. + + Rows before each request's length are unchanged. Modifies ``probs`` in + place and returns it; callers must own the storage. + """ + positions = torch.arange(probs.shape[1], device=probs.device) + invalid = positions.unsqueeze(0) >= verify_lens.to(positions.dtype).unsqueeze(1) + probs.masked_fill_(invalid.unsqueeze(-1), 0) + probs[..., 0].masked_fill_(invalid, 1) + return probs + + +def count_accepted_ragged( + *, + draft_tokens: torch.Tensor, + target_tokens: torch.Tensor, + verify_lens: torch.Tensor, +) -> torch.Tensor: + """Accepted draft-token count per request, honouring per-request windows. + + Positions at or beyond a request's ``verify_len`` hold stale padding and + must be masked to "no match" before the cumulative product, or a stale + slot can credit an acceptance the target never made. + + Args: + draft_tokens: ``[bs, max_len]`` padded drafted tokens. + target_tokens: ``[bs, max_len]`` padded target tokens. + verify_lens: ``[bs]`` positions actually verified per request. + """ + if draft_tokens.shape != target_tokens.shape: + raise ValueError( + f"draft {tuple(draft_tokens.shape)} and target " + f"{tuple(target_tokens.shape)} must have the same padded shape" + ) + max_len = draft_tokens.shape[1] + positions = torch.arange(max_len, device=draft_tokens.device) + valid = positions.unsqueeze(0) < verify_lens.to(positions.dtype).unsqueeze(1) + match = (draft_tokens == target_tokens) & valid + return torch.cumprod(match.int(), dim=-1).sum(1) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index dc068824b470..6aceef382b55 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -535,6 +535,29 @@ class SpecMetadata: # Normally, it equals 1 + runtime_draft_len. But for PARD, it equals 2 * runtime_draft_len. runtime_tokens_per_gen_step: int = 1 + # Ragged (per-request) verification split, set only when + # confidence-scheduled verification runs in ragged mode: + # + # verify_lens [num_gen] positions verified per generation request + # qo_indptr [num_gen + 1] exclusive prefix sum; request r owns + # [qo_indptr[r], qo_indptr[r + 1]) + # + # Both stay None on every other path. Built by + # ``_torch/speculative/dspark_ragged.RaggedVerifyLayout``. + verify_lens: Optional[torch.Tensor] = None + qo_indptr: Optional[torch.Tensor] = None + # Host-side ``sum(verify_lens)`` (reading it off the tensor would sync). + # Equals the captured token bucket after ``RaggedVerifyLayout.fill_bucket``. + total_verify_tokens: Optional[int] = None + + @property + def is_ragged_verify(self) -> bool: + """Whether this iteration uses per-request verify lengths. + + This property implies that generation rows use the packed layout. + """ + return self.verify_lens is not None + # Auto-detected per step from populated sampling params: # True if every request is greedy (no temp/top_k/top_p/min_p) and we can take # the argmax fast-path. False if any request needs sampling. @@ -1157,7 +1180,8 @@ def _scan_one_model_sampling( population. Does NOT allocate or fill GPU buffers, so it is safe to call before the CUDA graph key is built. """ - from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState + from tensorrt_llm._torch.pyexecutor.llm_request import ( + LlmRequestState, get_request_tokens_per_gen_step) from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import \ GREEDY_TEMPERATURE_THRESHOLD from tensorrt_llm.sampling_params import SamplingParams @@ -1219,6 +1243,8 @@ def _normalize_request_sampling_params( per_request_normalized: list[tuple[float, int, float, float, int]] = [] has_non_greedy_requests = False per_request_slot_ids: list[int] = [] + ragged_verify = self.is_ragged_verify + default_generation_tokens = 1 + self.runtime_draft_len for request in requests: sampling_config = request.sampling_config @@ -1227,8 +1253,16 @@ def _normalize_request_sampling_params( tp_val = sampling_config.top_p mp_val = sampling_config.min_p - # Context requests have no draft tokens yet. - num_tokens = 1 + self.runtime_draft_len if request.state == LlmRequestState.GENERATION_IN_PROGRESS else 1 + # Generation requests contribute one row per verified position, + # per-request under ragged verification: the buffers are indexed + # by flat token position, so a fixed stride would apply request + # r's sampling params to request r+1's tokens. + if request.state == LlmRequestState.GENERATION_IN_PROGRESS: + num_tokens = (get_request_tokens_per_gen_step( + request, default_generation_tokens) + if ragged_verify else default_generation_tokens) + else: + num_tokens = 1 ( temp_val, @@ -1548,6 +1582,34 @@ def commit_speculative_states( ... +def _padded_gen_draft_tokens(spec_metadata: SpecMetadata, num_gens: int, + runtime_draft_len: int) -> torch.Tensor: + """Unpack ragged-packed draft tokens into ``[num_gens, runtime_draft_len]``. + + The overlap path packs draft tokens by each request's own window, so the + buffer holds ``sum(verify_lens) - num_gens`` entries rather than a + ``num_gens x runtime_draft_len`` rectangle. Acceptance is naturally + rectangular, so scatter back out; positions past a request's window are + padding and get masked by ``count_accepted_ragged``. + """ + from .dspark_ragged import build_qo_indptr, scatter_ragged_to_padded + + # verify_lens counts tokens (bonus position + drafts); the draft buffer + # only holds the drafts. + verify_lens = spec_metadata.verify_lens + total_verify_tokens = spec_metadata.total_verify_tokens + if verify_lens is None or total_verify_tokens is None: + raise RuntimeError("ragged verification metadata is incomplete") + draft_lens = verify_lens - 1 + packed_len = total_verify_tokens - num_gens + return scatter_ragged_to_padded( + spec_metadata.draft_tokens[:packed_len], + verify_lens=draft_lens, + qo_indptr=build_qo_indptr(draft_lens), + max_len=runtime_draft_len, + ).to(torch.int32) + + class SpecWorkerBase(nn.Module, ABC): """ Base class for speculative decoding workers. @@ -1941,17 +2003,40 @@ def _sample_and_accept_draft_tokens_base( accepted_tokens[:num_contexts, 0] = target_tokens[:num_contexts] # Generation requests: verify draft tokens against target tokens - gen_target_tokens = target_tokens[num_contexts:].reshape( - num_gens, runtime_draft_len + 1) + if spec_metadata is not None and spec_metadata.is_ragged_verify: + from .dspark_ragged import (count_accepted_ragged, + scatter_ragged_to_padded) + verify_lens = spec_metadata.verify_lens + total = spec_metadata.total_verify_tokens + qo_indptr = spec_metadata.qo_indptr + if verify_lens is None or total is None or qo_indptr is None: + raise RuntimeError("ragged verification metadata is incomplete") + gen_target_tokens = scatter_ragged_to_padded( + target_tokens[num_contexts:num_contexts + total], + verify_lens=verify_lens, + qo_indptr=qo_indptr, + max_len=runtime_draft_len + 1) + # verify_lens counts tokens (bonus position included), so a + # request drafted verify_lens - 1 positions. Masking past that is + # what stops a stale padded slot from comparing equal and crediting + # an acceptance the target never made. + num_accepted_tokens[num_contexts:] += count_accepted_ragged( + draft_tokens=draft_tokens, + target_tokens=gen_target_tokens[:, :runtime_draft_len], + verify_lens=verify_lens - 1) + else: + gen_target_tokens = target_tokens[num_contexts:].reshape( + num_gens, runtime_draft_len + 1) + + # Compare draft tokens with target tokens using cumulative product + # Counts consecutive matches from the start + num_accepted_tokens[num_contexts:] += torch.cumprod(( + draft_tokens == gen_target_tokens[:, :runtime_draft_len]).int(), + dim=-1).sum(1) + accepted_tokens[num_contexts:, :runtime_draft_len + 1] = gen_target_tokens - # Compare draft tokens with target tokens using cumulative product - # Counts consecutive matches from the start - num_accepted_tokens[num_contexts:] += torch.cumprod( - (draft_tokens == gen_target_tokens[:, :runtime_draft_len]).int(), - dim=-1).sum(1) - # Apply force override if set num_accepted_tokens = self._apply_force_accepted_tokens( num_accepted_tokens, @@ -1959,6 +2044,14 @@ def _sample_and_accept_draft_tokens_base( runtime_draft_len, spec_metadata=spec_metadata) + if spec_metadata is not None and spec_metadata.is_ragged_verify: + # The synthetic-acceptance override caps at the batch-wide + # runtime_draft_len; under ragged a request may have verified + # fewer positions, and accepting past its window would commit + # tokens the target never produced. + num_accepted_tokens[num_contexts:] = torch.minimum( + num_accepted_tokens[num_contexts:], + spec_metadata.verify_lens.to(num_accepted_tokens.dtype)) return accepted_tokens, num_accepted_tokens def _apply_occurrence_penalties( @@ -1969,9 +2062,10 @@ def _apply_occurrence_penalties( No-op unless the deploy enabled the penalties and some request in the batch actually uses one. ``logits`` must already be in the normalized - ``[ctx (1 row), gen (draft_len + 1 rows)]`` layout, i.e. after - ``_reshape_logits_for_accept`` -- which is what makes PARD's wider raw - layout fit the same mapping. + layout after ``_reshape_logits_for_accept``: one row per context plus + either ``draft_len + 1`` rows per generation request (uniform) or each + request's packed ``verify_lens`` rows (ragged). This is also what makes + PARD's wider raw layout fit the same mapping. Returns the logits acceptance should read: a penalized copy when the penalties apply, otherwise the caller's tensor unchanged. @@ -1985,9 +2079,13 @@ def _apply_occurrence_penalties( # captured; whether it changes anything is decided on device by # ``penalty_active``, which the replayed kernel re-reads every step. draft_len = draft_tokens.shape[1] if draft_tokens.dim() > 1 else 0 - mapping = penalty_ops.build_row_mapping(spec_metadata, num_contexts, - batch_size, draft_len, - draft_tokens, logits.device) + mapping = penalty_ops.build_row_mapping(spec_metadata, + num_contexts, + batch_size, + draft_len, + draft_tokens, + logits.device, + num_logit_rows=logits.shape[0]) if mapping is None: return logits row_slots, intra_tokens, intra_valid = mapping @@ -2212,9 +2310,13 @@ def _compute_probs_from_logits( logits, temperatures, top_ks, top_ps) @staticmethod - def _zero_padding_rows(logits: torch.Tensor, spec_metadata, - num_contexts: int, batch_size: int, - rows_per_request: int) -> torch.Tensor: + def _zero_padding_rows(logits: torch.Tensor, + spec_metadata, + num_contexts: int, + batch_size: int, + rows_per_request: int, + *, + packed_verify: bool = False) -> torch.Tensor: """Zero the logits rows belonging to CUDA-graph padding requests. A padding request decodes from an uninitialized KV/hidden state, so its @@ -2231,20 +2333,31 @@ def _zero_padding_rows(logits: torch.Tensor, spec_metadata, ``prepare_penalty_buffers`` are the two places that publish it. Args: - logits: ``[(batch_size - num_contexts) * rows_per_request, vocab]``; - returned unchanged if the row count does not match. + logits: ``[num_generation_tokens, vocab]``. rows_per_request: logits rows each request contributes. + packed_verify: use the ragged verification lengths for target logits; + draft-step logits retain their uniform layout. Returns: ``logits`` with the padding requests' rows zeroed. """ slot_ids = getattr(spec_metadata, "batch_slot_ids", None) dummy_slot_row = getattr(spec_metadata, "dummy_slot_row", 0) num_gens = batch_size - num_contexts - if (slot_ids is None or dummy_slot_row <= 0 or num_gens <= 0 - or logits.shape[0] != num_gens * rows_per_request): + if slot_ids is None or dummy_slot_row <= 0 or num_gens <= 0: return logits is_padding = slot_ids[num_contexts:batch_size] == dummy_slot_row - if rows_per_request > 1: + if packed_verify and spec_metadata.is_ragged_verify: + lens = spec_metadata.verify_lens + if (lens is None or lens.dim() != 1 or lens.numel() != num_gens + or spec_metadata.total_verify_tokens != logits.shape[0]): + raise ValueError( + "ragged padding rows do not match the verification layout") + is_padding = torch.repeat_interleave(is_padding, + lens.to(torch.long), + output_size=logits.shape[0]) + elif logits.shape[0] != num_gens * rows_per_request: + return logits + elif rows_per_request > 1: is_padding = is_padding.repeat_interleave(rows_per_request) return logits.masked_fill(is_padding.unsqueeze(-1), 0.0) @@ -2324,6 +2437,12 @@ def _reshape_draft_tokens_for_accept(self, spec_metadata, num_gens, device): return torch.zeros((num_gens, runtime_draft_len), dtype=torch.int, device=device) + if spec_metadata.is_ragged_verify: + # Packed by per-request windows, not a rectangle; a plain reshape + # can happen to divide and silently attribute one request's drafts + # to another. + return _padded_gen_draft_tokens(spec_metadata, num_gens, + runtime_draft_len) return spec_metadata.draft_tokens.reshape(num_gens, runtime_draft_len) def _reshape_logits_for_accept(self, logits, num_contexts, num_gens, @@ -2380,9 +2499,21 @@ def _rejection_buffers_valid(self, draft_tokens, draft_len, stored_vocab, return False if draft_tokens.dim() != 2 or draft_tokens.shape[0] != num_gens: return False - # logits must cover context rows (1 each) + gen rows (draft_len + 1 each). + # logits must cover context rows (1 each) + however many gen rows the + # target was actually given: under ragged verification that is + # `total_verify_tokens`, not `num_gens * (draft_len + 1)`. logits_rows = logits.shape[0] if logits.dim() > 1 else 1 - if logits_rows < num_contexts + num_gens * (draft_len + 1): + if spec_metadata.is_ragged_verify: + # Fail closed on a half-built ragged layout: the branch below + # dereferences all three of these. + total_verify_tokens = spec_metadata.total_verify_tokens + if (total_verify_tokens is None or spec_metadata.verify_lens is None + or spec_metadata.qo_indptr is None): + return False + required_gen_rows = int(total_verify_tokens) + else: + required_gen_rows = num_gens * (draft_len + 1) + if logits_rows < num_contexts + required_gen_rows: return False # Slot ids for the gen subset must exist (range safety is guaranteed by # construction). @@ -2446,15 +2577,23 @@ def _sample_and_accept_draft_tokens_rejection( # === Generation subset: rejection sampling on the gen slice === if num_gens > 0: - num_gen_logits = num_gens * (runtime_draft_len + 1) + is_ragged = spec_metadata.is_ragged_verify + # The per-token sampling-parameter buffers are packed by the same + # per-request windows as the logits (see + # _scan_one_model_sampling), so one flat total covers both. + num_gen_logits = (spec_metadata.total_verify_tokens if is_ragged + else num_gens * (runtime_draft_len + 1)) gen_logits = logits[num_contexts:num_contexts + num_gen_logits] gen_start = num_contexts gen_end = num_contexts + num_gen_logits temperatures = spec_metadata.temperatures[gen_start:gen_end] - gen_logits = self._zero_padding_rows(gen_logits, spec_metadata, - num_contexts, batch_size, - runtime_draft_len + 1) + gen_logits = self._zero_padding_rows(gen_logits, + spec_metadata, + num_contexts, + batch_size, + runtime_draft_len + 1, + packed_verify=True) # The target distribution the acceptance test divides by. It must be # filtered exactly as the draft probs were (see advanced_sample_draft), # which is why both use the same backend selection on the same mode -- @@ -2464,9 +2603,19 @@ def _sample_and_accept_draft_tokens_rejection( spec_metadata.top_ks[gen_start:gen_end], spec_metadata.top_ps[gen_start:gen_end], spec_metadata.min_ps[gen_start:gen_end]) - target_probs = target_probs_flat.reshape(num_gens, - runtime_draft_len + 1, - vocab_size) + if is_ragged: + from .dspark_ragged import (fill_padded_rows_onehot, + scatter_ragged_to_padded) + target_probs = scatter_ragged_to_padded( + target_probs_flat, + verify_lens=spec_metadata.verify_lens, + qo_indptr=spec_metadata.qo_indptr, + max_len=runtime_draft_len + 1) + fill_padded_rows_onehot(target_probs, + verify_lens=spec_metadata.verify_lens) + else: + target_probs = target_probs_flat.reshape( + num_gens, runtime_draft_len + 1, vocab_size) draft_vocab_size = draft_probs.shape[-1] assert draft_probs.shape[0] == num_gens, ( @@ -2514,6 +2663,19 @@ def _sample_and_accept_draft_tokens_rejection( else: full_draft_probs = draft_probs + if is_ragged: + # Missing draft IDs are zero, so their proposal must be a + # point mass at target token 0. At the request-local bonus, + # acceptance emits 0 with probability q[0]; rejection samples + # q conditioned on a nonzero token. Together they sample q. + # Pad after d2t expansion, on private storage: target token 0 + # need not have a draft mapping and must not leak into reused + # buffers. + full_draft_probs = full_draft_probs.clone() + draft_lens = spec_metadata.verify_lens - 1 + fill_padded_rows_onehot(full_draft_probs, + verify_lens=draft_lens) + full_draft_tokens = draft_tokens.to(torch.int32).contiguous() # One entry per gen request; slot 0 of the step's offset window. @@ -2545,6 +2707,15 @@ def _sample_and_accept_draft_tokens_rejection( num_contexts, runtime_draft_len, spec_metadata=spec_metadata) + + if num_gens > 0 and spec_metadata.is_ragged_verify: + # The kernel walks the padded rectangle, so a request whose whole + # window was accepted keeps going into padding. verify_lens is the + # hard ceiling: accepting past it commits tokens the target never + # scored. + num_accepted_tokens[num_contexts:] = torch.minimum( + num_accepted_tokens[num_contexts:], + spec_metadata.verify_lens.to(num_accepted_tokens.dtype)) return accepted_tokens, num_accepted_tokens def _rng_state_per_request( diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 030fb18f4fd9..92f1305c1af2 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -130,6 +130,9 @@ class SampleStateTensorsSpec(SampleStateTensors): new_tokens_lens: torch.Tensor next_draft_tokens: torch.Tensor + # Actual token windows, aligned with SampleStateSpec.requests, when the + # worker selects verification lengths on device. + verify_lens: Optional[torch.Tensor] = None @dataclass(kw_only=True) @@ -145,6 +148,10 @@ class SampleStateSpec(SampleState): # instead would see the NEXT step's buffer, which update_requests itself # installs. draft_lens: Optional[list[int]] = None + # Per-request verify windows for this sampler step, keyed by request id. + # The overlap scheduler can stamp the next step on the live request before + # this state is consumed, so rewind accounting must use this snapshot. + verify_lens_snapshot: Optional[dict[int, int]] = None class SpecSampler(Sampler[SampleStateSpec], AsyncWorkerMixin): @@ -320,6 +327,57 @@ def _request_common_handling( request.py_draft_tokens = next_draft_tokens[request.py_seq_slot][:runtime_draft_len] request.py_decoding_iter += 1 + @staticmethod + def _verified_len( + request: LlmRequest, + runtime_draft_len: int, + verify_lens_snapshot: Optional[dict[int, int]], + token_window: Optional[int] = None, + ) -> int: + """Return the number of draft positions verified for one request.""" + if token_window is not None and token_window > 0: + return int(token_window) - 1 + verify_len = ( + verify_lens_snapshot.get(request.py_request_id) + if verify_lens_snapshot is not None + else None + ) + return runtime_draft_len if verify_len is None else int(verify_len) + + @staticmethod + def _snapshot_verify_lens( + requests: list[LlmRequest], + ) -> Optional[dict[int, int]]: + """Capture overlap-sensitive ragged windows for one sampler step.""" + snapshot = { + request.py_request_id: int(verify_len) + for request in requests + if (verify_len := getattr(request, "py_verify_len", None)) is not None + } + return snapshot or None + + @classmethod + def _snapshot_policy_windows_for_step( + cls, + requests: list[LlmRequest], + *, + native_uniform: bool, + host_snapshot_required: bool, + device_verify_lens_available: bool, + ) -> Optional[dict[int, int]]: + """Select one authoritative verification-window source for this step.""" + if type(native_uniform) is not bool or type(host_snapshot_required) is not bool: + raise ValueError("verification-window markers must be boolean") + if native_uniform and (host_snapshot_required or device_verify_lens_available): + raise ValueError("native-uniform verification cannot also publish ragged windows") + if host_snapshot_required and device_verify_lens_available: + raise ValueError("device windows and host-window snapshots are mutually exclusive") + if native_uniform: + return {} + if device_verify_lens_available: + return None + return cls._snapshot_verify_lens(requests) + def update_requests( self, state: SampleStateSpec, @@ -338,6 +396,9 @@ def update_requests( state.sampler_event.synchronize() new_tokens = state.host.new_tokens.tolist() new_tokens_lens_list = state.host.new_tokens_lens.tolist() + token_windows = ( + state.host.verify_lens.tolist() if state.host.verify_lens is not None else None + ) next_draft_tokens_list = state.host.next_draft_tokens.tolist() beam_idx = DEFAULT_BEAM_IDX runtime_draft_len = getattr(state, "runtime_draft_len", self.draft_len) @@ -365,16 +426,25 @@ def update_requests( # replaces py_draft_tokens with the next step's buffer, so its # length must not be used as the denominator (0 for the request's # prefill step, where nothing was verified). - req.py_num_draft_tokens_verified = ( - state.draft_lens[req_idx] if state.draft_lens is not None else 0 + verified_len = ( + runtime_draft_len + if token_windows is None and state.verify_lens_snapshot is None + else self._verified_len( + req, + runtime_draft_len, + state.verify_lens_snapshot, + token_windows[req_idx] if token_windows is not None else None, + ) ) - req.py_rewind_len = runtime_draft_len - req.py_num_accepted_draft_tokens + drafted_len = state.draft_lens[req_idx] if state.draft_lens is not None else 0 + req.py_num_draft_tokens_verified = min(drafted_len, verified_len) + req.py_rewind_len = verified_len - req.py_num_accepted_draft_tokens self._request_common_handling(req, next_draft_tokens_list, runtime_draft_len) def sample_async( self, scheduled_requests: ScheduledRequests, - outputs: dict[str, torch.Tensor], + outputs: dict[str, torch.Tensor | bool], num_context_logits_prefix_sum: list[int], ) -> SampleStateSpec: """ @@ -388,6 +458,9 @@ def sample_async( - new_tokens_lens: [batch] number of accepted tokens - next_draft_tokens: [batch, max_draft_len] draft tokens for next iter - next_new_tokens: [batch, max_draft_len + 1] input for next iter + - verify_lens: optional [batch] device-selected token windows + - native_uniform_verify: optional boolean declaring uniform windows + - host_policy_windows_snapshot: optional boolean declaring host windows num_context_logits_prefix_sum: Prefix sum of context logits (unused) Returns: @@ -424,6 +497,25 @@ def sample_async( o_next_new_tokens = outputs["next_new_tokens"][num_skip : num_skip + num_sampling_requests] runtime_draft_len = o_next_draft_tokens.shape[1] + o_verify_lens = outputs.get("verify_lens") + if o_verify_lens is not None: + if ( + not isinstance(o_verify_lens, torch.Tensor) + or o_verify_lens.dim() != 1 + or o_verify_lens.dtype not in (torch.int32, torch.int64) + or o_verify_lens.shape[0] < num_skip + num_sampling_requests + ): + raise ValueError( + "verify_lens must contain one integer token window per sampled request" + ) + o_verify_lens = o_verify_lens[num_skip : num_skip + num_sampling_requests] + verify_lens_snapshot = self._snapshot_policy_windows_for_step( + sampling_requests, + native_uniform=outputs.get("native_uniform_verify", False), + host_snapshot_required=outputs.get("host_policy_windows_snapshot", False), + device_verify_lens_available=o_verify_lens is not None, + ) + # Pad or truncate to match fixed-size store buffers for index_copy_. # The worker output width tracks runtime_draft_len, which dynamic draft # length shrinks below the statically allocated store width. @@ -466,6 +558,7 @@ def sample_async( new_tokens=self._copy_to_host(self.store.new_tokens), new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens), next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens), + verify_lens=(self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None), ) sampler_event = self._record_sampler_event() @@ -480,4 +573,5 @@ def sample_async( sampler_event=sampler_event, runtime_draft_len=runtime_draft_len, draft_lens=draft_lens, + verify_lens_snapshot=verify_lens_snapshot, ) 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 9784cc66fa9e..11d6d300e6f5 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 @@ -37,17 +37,23 @@ def _request( *, rewind: int = 0, accepted_draft_tokens: int = 0, + draft_tokens: list[int] | None = None, + verify_len: int | None = None, complete: bool = False, ) -> SimpleNamespace: - return SimpleNamespace( + request = SimpleNamespace( py_request_id=request_id, py_rewind_len=rewind, py_num_accepted_draft_tokens=accepted_draft_tokens, + py_draft_tokens=draft_tokens, max_beam_num_tokens=201, state=LlmRequestState.GENERATION_COMPLETE if complete else LlmRequestState.GENERATION_IN_PROGRESS, ) + if verify_len is not None: + request.py_verify_len = verify_len + return request def _cache(*, capacity: int = 256, active: bool = True) -> MagicMock: @@ -139,6 +145,37 @@ def test_dynamic_tree_reserved_capacity(is_draft: bool, expected_capacity: int) cache.resize.assert_called_once_with(expected_capacity, 200) +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("verify_len", "rewind", "accepted", "expected_capacity"), + [(2, 1, 1, 252), (5, 2, 3, 254)], + ids=["reclaims-unverified-suffix", "full-window-keeps-uniform-accounting"], +) +def test_target_reclaims_only_ragged_unverified_capacity( + verify_len: int, rewind: int, accepted: int, expected_capacity: int +) -> None: + manager = _manager(is_draft=False, kv_reserve_draft_tokens=5) + request = _request( + 1, + rewind=rewind, + accepted_draft_tokens=accepted, + draft_tokens=[1] * 5, + verify_len=verify_len, + ) + cache = _cache() + manager.kv_cache_map[request.py_request_id] = cache + + # The reservation belongs to the completed step; live request fields may + # already describe the next overlap iteration. + manager._allocated_draft_lens[request.py_request_id] = 5 + request.py_verify_len = 1 + request.py_draft_tokens = [9] + manager.update_resources(SimpleNamespace(generation_requests=[request])) + + cache.resize.assert_called_once_with(expected_capacity, 200) + assert request.py_request_id not in manager._allocated_draft_lens + + def test_capacity_only_completion_preserves_history() -> None: manager = _manager(is_draft=False, kv_compression_manages_history=True) request = _request(1, complete=True) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py new file mode 100644 index 000000000000..182cc5b0c221 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py @@ -0,0 +1,500 @@ +# 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 ragged verification layout and packing tests.""" + +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark_ragged import ( + RaggedVerifyLayout, + build_qo_indptr, + build_row_maps_device, + count_accepted_ragged, + fill_bucket_device, + fill_padded_rows_onehot, + resolve_ragged_pad_split, + row_ids_from_lens, + scatter_ragged_to_padded, +) + +BLOCK = 7 +pytestmark = pytest.mark.cpu_only + + +def _layout(lens, bucket=None): + t = torch.tensor(lens, dtype=torch.int32) + bucket = sum(lens) if bucket is None else bucket + return RaggedVerifyLayout.from_verify_lens( + t, graph_num_tokens=bucket, total_verify_tokens=sum(lens) + ) + + +# -------------------------------------------------------------------------- +# packing +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "lens, expected", + [ + ([3, 1, 4, 2], [0, 3, 4, 8, 10]), + ([], [0]), # empty batch still yields the leading 0 + ], +) +def test_qo_indptr_is_the_exclusive_prefix_sum(lens, expected): + assert build_qo_indptr(torch.tensor(lens, dtype=torch.int32)).tolist() == expected + + +def test_build_qo_indptr_rejects_2d(): + with pytest.raises(ValueError, match="1-D"): + build_qo_indptr(torch.ones(2, 3, dtype=torch.int32)) + + +# -------------------------------------------------------------------------- +# padding reclamation +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "lens, bucket", + [ + ([2, 1, 1], 12), # 4 used, 8 spare: padding becomes real verification + ([3, 2, 3], 8), # already exact: the fill is a no-op + ], +) +def test_fill_bucket_converts_padding_into_real_verification(lens, bucket): + """The step already pays for the bucket -- spare slots must not be wasted.""" + lay = _layout(lens, bucket=bucket) + filled = lay.fill_bucket(max_verify_len=BLOCK) + assert filled.total_verify_tokens == bucket + assert int(filled.verify_lens.sum()) == bucket + filled.validate(max_verify_len=BLOCK, exact_fill=True) + + +def test_fill_bucket_spreads_over_real_requests_first(): + """Round-robin keeps the extra where survival is still plausible.""" + lay = _layout([1, 1, 1, 1], bucket=8) + filled = lay.fill_bucket(max_verify_len=BLOCK) + lens = filled.verify_lens.tolist() + assert sum(lens) == 8 + assert max(lens) - min(lens) <= 1, f"unbalanced: {lens}" + + +def test_fill_bucket_adds_pad_rows_to_reach_the_captured_batch_size(): + """Row count must hit padded_bs; slack goes to real rows before pad rows.""" + lay = _layout([3, 2], bucket=12) + filled = lay.fill_bucket(max_verify_len=BLOCK, padded_bs=4) + assert filled.bs == 4 + assert filled.num_real_requests == 2 + assert filled.num_pad_requests == 2 + assert int(filled.verify_lens.sum()) == 12 + filled.validate(max_verify_len=BLOCK, exact_fill=True) + # Real rows absorb all the slack; the pad rows stay at the 1-token minimum + # (a pad row with 0 tokens would break qo_indptr slicing). + lens = filled.verify_lens.tolist() + assert lens[2:] == [1, 1] + assert lens[0] + lens[1] == 10 + + +def test_fill_bucket_raises_when_the_bucket_cannot_be_reached(): + """Silently returning a short batch would desync attention from the MoE.""" + lay = _layout([1, 1], bucket=100) # capacity 2*7 = 14 < 100 + with pytest.raises(ValueError, match="exceeds what .* rows can absorb"): + lay.fill_bucket(max_verify_len=BLOCK) + + +@pytest.mark.parametrize( + "lens, bucket, padded_bs", + [ + ([5, 5], 8, None), # 10 real tokens already > 8 + # 8 real tokens + 6 one-token pad rows needs 14, not 10: the pad-row + # floor counts against the bucket too. A short batch would desync + # attention from the MoE. + ([4, 4], 10, 8), + ], +) +def test_fill_bucket_raises_when_the_bucket_is_too_small(lens, bucket, padded_bs): + lay = _layout(lens, bucket=bucket) + with pytest.raises(ValueError, match="too small"): + lay.fill_bucket(max_verify_len=BLOCK, padded_bs=padded_bs) + + +def test_fill_bucket_rejects_a_padded_bs_below_the_real_count(): + lay = _layout([2, 2, 2], bucket=8) + with pytest.raises(ValueError, match="cannot hold this batch"): + lay.fill_bucket(max_verify_len=BLOCK, padded_bs=2) + + +def test_validate_exact_fill_catches_a_short_batch(): + """The check that would otherwise only surface as an attention/MoE mismatch.""" + lay = _layout([2, 2], bucket=8) # 4 != 8 + lay.validate(max_verify_len=BLOCK) # non-exact form still passes + with pytest.raises(ValueError, match="do not fill the captured bucket"): + lay.validate(max_verify_len=BLOCK, exact_fill=True) + + +def test_fill_bucket_needs_the_host_total(): + lay = RaggedVerifyLayout.from_verify_lens( + torch.tensor([2, 2], dtype=torch.int32), graph_num_tokens=8 + ) + with pytest.raises(ValueError, match="host-side total"): + lay.fill_bucket(max_verify_len=BLOCK) + + +# -------------------------------------------------------------------------- +# invariants +# -------------------------------------------------------------------------- + + +def test_validate_rejects_a_request_that_verifies_nothing(): + """verify_len 0 means the request makes no progress -- it would stall.""" + lay = _layout([2, 0, 3], bucket=8) + with pytest.raises(ValueError, match="anchor token"): + lay.validate() + + +def test_validate_rejects_overflowing_the_bucket(): + lay = _layout([5, 5], bucket=8) + with pytest.raises(ValueError, match="exceed the captured bucket"): + lay.validate() + + +def test_validate_rejects_a_stale_indptr(): + lay = _layout([2, 3], bucket=8) + lay.qo_indptr = torch.tensor([0, 1, 5], dtype=torch.int32) + with pytest.raises(ValueError, match="not the prefix sum"): + lay.validate() + + +# -------------------------------------------------------------------------- +# packed <-> padded, and ragged acceptance counting +# -------------------------------------------------------------------------- + + +def test_row_ids_map_each_packed_token_to_its_request(): + lens = torch.tensor([3, 1, 2], dtype=torch.int32) + assert row_ids_from_lens(lens, total=6).tolist() == [0, 0, 0, 1, 2, 2] + + +def test_row_ids_requires_the_host_total(): + """``total`` keeps the expansion sync-free -- without it torch reads the + cumsum back to the host, which is illegal inside the captured graph.""" + lens = torch.tensor([3, 1, 2], dtype=torch.int32) + with pytest.raises(TypeError): + row_ids_from_lens(lens) + + +def test_scatter_unpacks_a_flat_batch_into_rows(): + lens = torch.tensor([3, 1, 2], dtype=torch.int32) + lay = _layout([3, 1, 2], bucket=8) + flat = torch.tensor([10, 11, 12, 20, 30, 31]) + out = scatter_ragged_to_padded( + flat, verify_lens=lens, qo_indptr=lay.qo_indptr, max_len=4, pad_value=-1 + ) + assert out.tolist() == [[10, 11, 12, -1], [20, -1, -1, -1], [30, 31, -1, -1]] + + +@pytest.mark.parametrize( + "draft, target, lens, expected", + [ + # the run of matches stops at the first mismatch + ([[1, 2, 3, 4]], [[1, 2, 9, 4]], [4], [2]), + # a request only gets credit inside its window, even when the stale + # padding beyond it happens to match -- unmasked padding would credit + # an acceptance the target never made (wrong output, no error). + ([[1, 2, 3, 4], [5, 6, 7, 8]], [[1, 2, 3, 4], [5, 6, 7, 8]], [4, 2], [4, 2]), + ], +) +def test_accept_count_respects_each_request_window(draft, target, lens, expected): + got = count_accepted_ragged( + draft_tokens=torch.tensor(draft), + target_tokens=torch.tensor(target), + verify_lens=torch.tensor(lens), + ) + assert got.tolist() == expected + + +def test_accept_count_matches_the_uniform_path_when_lengths_are_equal(): + """Equal lengths must reproduce the existing rectangular cumprod exactly.""" + torch.manual_seed(9) + bs, L = 6, 5 + draft = torch.randint(0, 4, (bs, L)) + target = torch.randint(0, 4, (bs, L)) + lens = torch.full((bs,), L) + ragged = count_accepted_ragged(draft_tokens=draft, target_tokens=target, verify_lens=lens) + uniform = torch.cumprod((draft == target).int(), dim=-1).sum(1) + assert torch.equal(ragged, uniform) + + +def test_accept_count_rejects_mismatched_shapes(): + with pytest.raises(ValueError, match="same padded shape"): + count_accepted_ragged( + draft_tokens=torch.zeros(2, 3, dtype=torch.long), + target_tokens=torch.zeros(2, 4, dtype=torch.long), + verify_lens=torch.tensor([3, 3]), + ) + + +# --------------------------------------------------------------------------- +# fill_padded_rows_onehot +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("padding_value", [0.0, 0.2, float("nan")]) +def test_padding_rows_become_a_valid_distribution(padding_value): + probs = torch.full((2, 3, 5), padding_value) + probs[0, :2] = torch.tensor([0.2, 0.3, 0.1, 0.2, 0.2]) + probs[1, :1] = torch.tensor([0.5, 0.1, 0.1, 0.2, 0.1]) + fill_padded_rows_onehot(probs, verify_lens=torch.tensor([2, 1])) + + # Real positions untouched. + assert torch.allclose(probs[0, 0], torch.tensor([0.2, 0.3, 0.1, 0.2, 0.2])) + # Clear stale mass (including non-finite values), not just the token-0 entry. + assert torch.allclose(probs[0, 2].sum(), torch.tensor(1.0)) + assert torch.allclose(probs[1, 1].sum(), torch.tensor(1.0)) + assert torch.allclose(probs[1, 2].sum(), torch.tensor(1.0)) + assert torch.equal(probs[0, 2], torch.tensor([1.0, 0, 0, 0, 0])) + + # Full-length rows are left entirely alone (guards the >= boundary). + full = torch.full((2, 3, 4), 0.25) + before = full.clone() + fill_padded_rows_onehot(full, verify_lens=torch.tensor([3, 3])) + assert torch.equal(full, before) + + +# --------------------------------------------------------------------------- +# Bucket fitting: the runtime batch must land on a captured token total +# --------------------------------------------------------------------------- + + +def _fit(real_token_lens, bucket, padded_bs, max_verify_len): + """Mirror of ModelEngine.fit_ragged_verify_lens' core.""" + lens = torch.tensor(real_token_lens, dtype=torch.int32) + layout = RaggedVerifyLayout.from_verify_lens( + lens, graph_num_tokens=bucket, total_verify_tokens=sum(real_token_lens) + ) + return layout.fill_bucket(max_verify_len=max_verify_len, padded_bs=padded_bs) + + +def test_fit_never_exceeds_the_per_request_ceiling(): + got = _fit([1, 1, 1], bucket=12, padded_bs=4, max_verify_len=3) + assert max(got.verify_lens.tolist()) <= 3 + + +@pytest.mark.parametrize( + "bucket, expected_pad_len", + [ + # 3 real rows capped at 6 hold up to 18; bucket 18 minus 5 one-token + # pad rows leaves 13, which the real rows absorb -- pads stay minimal. + (18, 1), + # The case that used to have no solution: bucket 24 with 5 one-token + # pad rows would demand 19 from real rows that cap at 18, so the pad + # rows must absorb what the real rows cannot. + (24, 2), + ], +) +def test_pad_rows_grow_when_the_real_rows_cannot_absorb_the_bucket(bucket, expected_pad_len): + padded_bs, max_verify_len = 8, 6 + real = [2, 3, 2] + n_pad = padded_bs - len(real) + split = resolve_ragged_pad_split( + bucket=bucket, + num_real_requests=len(real), + total_real_tokens=sum(real), + padded_bs=padded_bs, + max_verify_len=max_verify_len, + ) + assert split is not None + assert split.pad_len == expected_pad_len + + real_target = split.real_target + assert sum(real) <= real_target <= len(real) * max_verify_len + got = _fit(real, bucket=real_target, padded_bs=len(real), max_verify_len=max_verify_len) + # The total the forward sees must be exactly the captured bucket. + assert int(got.verify_lens.sum()) + n_pad * split.pad_len == bucket + + +def test_pad_length_never_yields_an_infeasible_real_target(): + """Over every shape the scheduler can produce, the derived split must either + reconstruct the bucket exactly or decline outright (never a missed bucket).""" + max_verify_len = 6 + checked = declined = 0 + for padded_bs in (1, 2, 4, 8, 16, 32, 64): + grid = sorted({padded_bs * (t + 1) for t in range(max_verify_len)}) + for n_real in range(1, padded_bs + 1): + n_pad = padded_bs - n_real + for floor in range(n_real, n_real * max_verify_len + 1): + # The scheduler rounds up from "what the batch already needs + # plus one token per pad row". + fits = [b for b in grid if b >= floor + n_pad] + if not fits: + continue + bucket = fits[0] + checked += 1 + split = resolve_ragged_pad_split( + bucket=bucket, + num_real_requests=n_real, + total_real_tokens=floor, + padded_bs=padded_bs, + max_verify_len=max_verify_len, + ) + if split is None: + declined += 1 + continue + real_target = split.real_target + assert floor <= real_target <= n_real * max_verify_len, ( + f"padded_bs={padded_bs} n_real={n_real} bucket={bucket} " + f"floor={floor} pad_len={split.pad_len} -> {real_target}" + ) + assert real_target + n_pad * split.pad_len == bucket + + assert checked > 10000, f"sweep too small to be meaningful: {checked}" + # Declining is safe (the step stays uniform) but should stay rare, or the + # bucket grid is too coarse to be worth capturing. + assert declined / checked < 0.15, f"{declined}/{checked} shapes have no valid split" + + +# --------------------------------------------------------------------------- +# device fill / row-map parity: fill_bucket_device and build_row_maps_device +# must reproduce the host path token-for-token. +# --------------------------------------------------------------------------- + + +def _host_fill(lens, padded_bs, bucket, max_verify_len): + layout = RaggedVerifyLayout.from_verify_lens( + torch.tensor(lens, dtype=torch.int32), + graph_num_tokens=bucket, + total_verify_tokens=sum(lens), + ) + return layout.fill_bucket(max_verify_len=max_verify_len, padded_bs=padded_bs) + + +def _device_fill(lens, padded_bs, bucket, max_verify_len): + padded = torch.ones(padded_bs, dtype=torch.int32) + padded[: len(lens)] = torch.tensor(lens, dtype=torch.int32) + return fill_bucket_device( + padded, + num_real=torch.tensor(len(lens)), + graph_num_tokens=bucket, + max_verify_len=max_verify_len, + ) + + +def _feasible(lens, padded_bs, bucket, max_verify_len): + n_pad = padded_bs - len(lens) + return sum(lens) + n_pad <= bucket <= padded_bs * max_verify_len + + +def _all_lens(n_real, max_verify_len): + """Every [1, max]^n_real length vector.""" + if n_real == 0: + yield [] + return + for head in range(1, max_verify_len + 1): + for tail in _all_lens(n_real - 1, max_verify_len): + yield [head] + tail + + +@pytest.mark.parametrize("max_verify_len", [2, 4, 6]) +def test_fill_parity_exhaustive_small(max_verify_len): + """Exact host/device agreement over every feasible small case.""" + checked = 0 + for n_real in (1, 2, 3): + for padded_bs in (n_real, n_real + 1, n_real + 3): + for lens in _all_lens(n_real, max_verify_len): + for bucket in range(padded_bs, padded_bs * max_verify_len + 1): + if not _feasible(lens, padded_bs, bucket, max_verify_len): + continue + host = _host_fill(lens, padded_bs, bucket, max_verify_len) + dev = _device_fill(lens, padded_bs, bucket, max_verify_len) + assert torch.equal(dev, host.verify_lens), ( + f"lens={lens} padded_bs={padded_bs} bucket={bucket} " + f"max={max_verify_len}: device {dev.tolist()} != " + f"host {host.verify_lens.tolist()}" + ) + checked += 1 + assert checked > 100 + + +def test_fill_parity_randomized_large(): + """Production-shaped cases: bs up to 128, max_verify_len 6.""" + gen = torch.Generator().manual_seed(20260806) + max_verify_len = 6 + for _ in range(200): + n_real = int(torch.randint(1, 129, (1,), generator=gen)) + padded_bs = n_real + int(torch.randint(0, 9, (1,), generator=gen)) + lens = torch.randint(1, max_verify_len + 1, (n_real,), generator=gen).tolist() + lo, hi = sum(lens) + (padded_bs - n_real), padded_bs * max_verify_len + bucket = int(torch.randint(lo, hi + 1, (1,), generator=gen)) + host = _host_fill(lens, padded_bs, bucket, max_verify_len) + dev = _device_fill(lens, padded_bs, bucket, max_verify_len) + assert torch.equal(dev, host.verify_lens) + assert int(dev.sum()) == bucket + + +def test_row_maps_match_prepare_semantics(): + """req_idx/correction must equal what prepare() stages: token o of a window + v gets correction o - v + 1, so the extent walks kv_len - v + 1 .. kv_len.""" + gen = torch.Generator().manual_seed(7) + for _ in range(50): + bs = int(torch.randint(1, 65, (1,), generator=gen)) + lens = torch.randint(1, 7, (bs,), generator=gen).to(torch.int32) + total = int(lens.sum()) + req_idx, corr = build_row_maps_device(lens, graph_num_tokens=total) + want_req, want_corr = [], [] + for r, v in enumerate(lens.tolist()): + for o in range(v): + want_req.append(r) + want_corr.append(o - v + 1) + assert req_idx.tolist() == want_req + assert corr.tolist() == want_corr + # Composing with a kv_lens gather is refresh_ragged_row_kv_lens: + # the extents must end exactly at each request's kv_len. + kv_lens = torch.randint(10, 1000, (bs,), generator=gen).to(torch.int32) + extents = kv_lens[req_idx] + corr + ends = torch.cumsum(lens.to(torch.long), 0) - 1 + assert torch.equal(extents[ends], kv_lens) + assert int(extents.min()) >= 1 + + +def test_fill_pad_fill_constrained_split(): + """With pad_fill given, every pad row carries exactly that many tokens, the + real rows absorb all slack, and the total still hits the bucket.""" + gen = torch.Generator().manual_seed(3) + max_verify_len = 6 + for _ in range(100): + n_real = int(torch.randint(1, 65, (1,), generator=gen)) + n_pad = int(torch.randint(0, 9, (1,), generator=gen)) + padded_bs = n_real + n_pad + pad_fill = int(torch.randint(1, max_verify_len + 1, (1,), generator=gen)) + lens = torch.randint(2, max_verify_len + 1, (n_real,), generator=gen) + lo = int(lens.sum()) + n_pad * pad_fill + hi = n_real * max_verify_len + n_pad * pad_fill + if lo > hi: + continue + bucket = int(torch.randint(lo, hi + 1, (1,), generator=gen)) + padded = torch.ones(padded_bs, dtype=torch.int32) + padded[:n_real] = lens.to(torch.int32) + filled = fill_bucket_device( + padded, + num_real=torch.tensor(n_real), + graph_num_tokens=bucket, + max_verify_len=max_verify_len, + pad_fill=pad_fill, + ) + assert int(filled.sum()) == bucket + assert (filled[n_real:] == pad_fill).all() + assert (filled[:n_real] >= lens.to(torch.int32)).all() + assert int(filled[:n_real].max()) <= max_verify_len diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged_metadata.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged_metadata.py new file mode 100644 index 000000000000..89b294f36073 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged_metadata.py @@ -0,0 +1,90 @@ +# 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. +"""Per-request token counts and ragged gather ordering.""" + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.llm_request import get_request_tokens_per_gen_step +from tensorrt_llm._torch.speculative.dspark_ragged import ragged_gather_index_lists + +pytestmark = pytest.mark.cpu_only + + +class _Req: + """Minimal stand-in exposing only what the layout code reads.""" + + def __init__(self, verify_len=None): + if verify_len is not None: + self.py_verify_len = verify_len + + +_MISSING = object() + + +@pytest.mark.parametrize( + "verify_len, runtime_tokens, expected", + [ + # Every non-DSpark request lands here: the attribute does not exist. + (_MISSING, 4, 4), + # The ragged path declined this request: the attribute is None. + (None, 4, 4), + # py_verify_len counts draft positions, matching runtime_draft_len; + # the accepted token makes it one more. + (2, 6, 3), + # A zero-draft window still carries the accepted token. + (0, 6, 1), + ], +) +def test_tokens_per_gen_step_resolves_each_request_window(verify_len, runtime_tokens, expected): + req = _Req() + if verify_len is not _MISSING: + req.py_verify_len = verify_len + assert get_request_tokens_per_gen_step(req, runtime_tokens) == expected + + +def test_uniform_counts_reproduce_the_strided_gather(): + # This is the safety property the whole change rests on: with equal counts + # the index list picks exactly the elements tensor[slots, :width] does. + table = torch.arange(6 * 5).reshape(6, 5) + slots = [3, 0, 4] + width = 5 + + rows, cols = ragged_gather_index_lists(slots, [width] * len(slots)) + ragged = table[torch.tensor(rows), torch.tensor(cols)] + strided = table[torch.tensor(slots), :width].flatten() + + assert torch.equal(ragged, strided) + + +def test_ragged_counts_pack_each_request_window(): + table = torch.arange(4 * 5).reshape(4, 5) + rows, cols = ragged_gather_index_lists([2, 0], [3, 1]) + gathered = table[torch.tensor(rows), torch.tensor(cols)] + + # slot 2 columns 0..2, then slot 0 column 0 + assert gathered.tolist() == [10, 11, 12, 0] + + # A zero count contributes nothing while its neighbours keep their order, + # and the empty batch degenerates cleanly. + assert ragged_gather_index_lists([1, 2, 3], [2, 0, 1]) == ([1, 1, 3], [0, 1, 0]) + assert ragged_gather_index_lists([], []) == ([], []) + + +def test_mismatched_lengths_raise(): + with pytest.raises(ValueError, match="same length"): + ragged_gather_index_lists([0, 1], [3]) + with pytest.raises(ValueError, match="negative gather count"): + ragged_gather_index_lists([0], [-1]) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py new file mode 100644 index 000000000000..45fa4b430241 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py @@ -0,0 +1,714 @@ +# 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. +"""Ragged one-model speculative sampling and penalty tests.""" + +import types +from typing import Optional + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.pyexecutor.sampler import penalties as penalty_ops +from tensorrt_llm._torch.speculative import interface as interface_ops +from tensorrt_llm._torch.speculative import spec_sampler_base as sampler_ops +from tensorrt_llm._torch.speculative.interface import SpecMetadata, SpecWorkerBase +from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler + +pytestmark = pytest.mark.cpu_only + + +class _StubSpecWorker(SpecWorkerBase): + """Concrete ``SpecWorkerBase`` that stubs out the abstract API.""" + + @property + def max_draft_len(self) -> int: + return 8 + + def _forward_impl(self, *args: object, **kwargs: object) -> None: + raise NotImplementedError + + +def _make_worker(value: Optional[float] = None) -> _StubSpecWorker: + # These tests cover CPU layout/acceptance, not native backend initialization. + worker = _StubSpecWorker.__new__(_StubSpecWorker) + torch.nn.Module.__init__(worker) + worker.force_num_accepted_tokens = 0.0 if value is None else value + worker._d2t = None + return worker + + +def _penalty_mapping_meta( + slot_ids: list[int], verify_lens: Optional[list[int]] = None +) -> types.SimpleNamespace: + metadata = types.SimpleNamespace( + batch_slot_ids=torch.tensor(slot_ids, dtype=torch.int64), + is_ragged_verify=verify_lens is not None, + verify_lens=None, + qo_indptr=None, + total_verify_tokens=None, + ) + if verify_lens is not None: + lens = torch.tensor(verify_lens, dtype=torch.int32) + metadata.verify_lens = lens + metadata.qo_indptr = torch.cat( + [torch.zeros(1, dtype=torch.int32), torch.cumsum(lens, dim=0)] + ) + metadata.total_verify_tokens = sum(verify_lens) + return metadata + + +def test_occurrence_penalty_uniform_row_mapping_is_unchanged(): + """An all-full ragged layout must reproduce the existing uniform mapping.""" + draft_len = 4 + draft_tokens = torch.tensor([[11, 12, 13, 14], [21, 22, 23, 24]]) + uniform = _penalty_mapping_meta([3, 7, 9]) + ragged = _penalty_mapping_meta([3, 7, 9], [draft_len + 1, draft_len + 1]) + num_rows = 1 + 2 * (draft_len + 1) + + uniform_mapping = penalty_ops.build_row_mapping( + uniform, + num_contexts=1, + batch_size=3, + draft_len=draft_len, + draft_tokens=draft_tokens, + device=torch.device("cpu"), + num_logit_rows=num_rows, + ) + ragged_mapping = penalty_ops.build_row_mapping( + ragged, + num_contexts=1, + batch_size=3, + draft_len=draft_len, + draft_tokens=draft_tokens, + device=torch.device("cpu"), + num_logit_rows=num_rows, + ) + + assert uniform_mapping is not None and ragged_mapping is not None + for uniform_tensor, ragged_tensor in zip(uniform_mapping, ragged_mapping): + assert torch.equal(uniform_tensor, ragged_tensor) + + +def test_occurrence_penalty_ragged_row_mapping_uses_each_request_window(): + """Packed rows must stay with their request and see only its earlier drafts.""" + draft_tokens = torch.tensor([[11, 12, 13, 14], [21, 22, 23, 24]]) + metadata = _penalty_mapping_meta([3, 7, 9], [2, 4]) + + mapping = penalty_ops.build_row_mapping( + metadata, + num_contexts=1, + batch_size=3, + draft_len=4, + draft_tokens=draft_tokens, + device=torch.device("cpu"), + num_logit_rows=7, + ) + + assert mapping is not None + row_slots, intra_tokens, intra_valid = mapping + assert row_slots.tolist() == [3, 7, 7, 9, 9, 9, 9] + assert intra_tokens.tolist() == [ + [0, 0, 0, 0], + [11, 12, 13, 14], + [11, 12, 13, 14], + [21, 22, 23, 24], + [21, 22, 23, 24], + [21, 22, 23, 24], + [21, 22, 23, 24], + ] + assert intra_valid.tolist() == [ + [False, False, False, False], + [False, False, False, False], + [True, False, False, False], + [False, False, False, False], + [True, False, False, False], + [True, True, False, False], + [True, True, True, False], + ] + + +def test_occurrence_penalty_applies_to_ragged_packed_rows(monkeypatch): + """The worker must run, rather than silently skip, the ragged penalty pass.""" + metadata = _penalty_mapping_meta([3, 7, 9], [2, 4]) + metadata.enable_penalty = True + captured = {} + + def _apply(logits, spec_metadata, row_slots, intra_tokens, intra_valid): + captured["row_slots"] = row_slots.clone() + captured["intra_tokens"] = intra_tokens.clone() + captured["intra_valid"] = intra_valid.clone() + logits[:, 0].copy_(row_slots.to(logits.dtype)) + + monkeypatch.setattr(penalty_ops, "apply_penalties", _apply) + logits = torch.zeros((7, 3)) + draft_tokens = torch.tensor([[11, 12, 13, 14], [21, 22, 23, 24]]) + + penalized = _make_worker()._apply_occurrence_penalties( + logits, draft_tokens, num_contexts=1, batch_size=3, spec_metadata=metadata + ) + + assert torch.equal(logits, torch.zeros_like(logits)) + assert penalized[:, 0].tolist() == [3.0, 7.0, 7.0, 9.0, 9.0, 9.0, 9.0] + assert captured["row_slots"].numel() == logits.shape[0] + assert captured["intra_valid"].sum().item() == 7 + + +def test_verify_window_snapshot_wins_over_next_overlap_step(): + request = types.SimpleNamespace(py_request_id=7, py_verify_len=2) + snapshot = SpecSampler._snapshot_verify_lens([request]) + request.py_verify_len = 5 + + assert snapshot == {7: 2} + assert SpecSampler._verified_len(request, 5, snapshot) == 2 + + +def test_uniform_verify_window_keeps_runtime_draft_length(): + request = types.SimpleNamespace(py_request_id=7) + + snapshot = SpecSampler._snapshot_verify_lens([request]) + request.py_verify_len = 1 + assert snapshot is None + assert SpecSampler._verified_len(request, 5, snapshot) == 5 + + +@pytest.mark.parametrize("force_accept", [0.0, 1.0, 3.0]) +def test_ragged_strict_acceptance_stops_at_each_request_window(force_accept): + worker = _make_worker(force_accept) + worker._sample_tokens_for_batch = lambda *args: torch.tensor( + [1, 100, 2, 8, 4, 101], dtype=torch.int64 + ) + metadata = _penalty_mapping_meta([0, 1], [2, 4]) + metadata.is_cuda_graph = False + draft_tokens = torch.tensor([[1, 9, 9], [2, 3, 4]], dtype=torch.int64) + + _, accepted = worker._sample_and_accept_draft_tokens_base( + logits=torch.zeros((6, 8)), + draft_tokens=draft_tokens, + num_contexts=0, + batch_size=2, + spec_metadata=metadata, + ) + + assert accepted.tolist() == ([2, 4] if force_accept == 3.0 else [2, 2]) + + +@pytest.mark.parametrize( + "windows, packed_drafts, expected", + [ + ([1, 3], [21, 22], [[0, 0], [21, 22]]), + ([1, 1], [], [[0, 0], [0, 0]]), + ], +) +def test_packed_drafts_allow_anchor_only_requests(windows, packed_drafts, expected): + metadata = _penalty_mapping_meta([0, 1], windows) + # The trailing capacity is stale and must not become another request's draft. + metadata.draft_tokens = torch.tensor([*packed_drafts, 99, 99], dtype=torch.int32) + + padded = interface_ops._padded_gen_draft_tokens(metadata, num_gens=2, runtime_draft_len=2) + + assert padded.dtype == torch.int32 + assert padded.tolist() == expected + + +def test_ragged_rejection_guard_uses_packed_logit_count(): + metadata = types.SimpleNamespace( + draft_probs=torch.empty((2, 3, 8)), + batch_slot_ids=torch.arange(2, dtype=torch.long), + is_ragged_verify=True, + verify_lens=torch.tensor([2, 4], dtype=torch.int32), + qo_indptr=torch.tensor([0, 2, 6], dtype=torch.int32), + total_verify_tokens=6, + ) + draft_tokens = torch.zeros((2, 3), dtype=torch.int64) + valid_logits = torch.zeros((6, 8)) + short_logits = torch.zeros((5, 8)) + + assert SpecWorkerBase._rejection_buffers_valid( + object(), draft_tokens, 3, 8, 0, 2, valid_logits, metadata + ) + assert not SpecWorkerBase._rejection_buffers_valid( + object(), draft_tokens, 3, 8, 0, 2, short_logits, metadata + ) + + +@pytest.mark.parametrize("packed_verify", [False, True]) +def test_padding_mask_distinguishes_target_windows_from_draft_rows(packed_verify): + metadata = _penalty_mapping_meta([8, 3, 9], [2, 4]) + metadata.dummy_slot_row = 9 + logits = torch.tensor([[1.0, 2.0], [float("nan"), float("inf")]]) + expected = torch.tensor([[1.0, 2.0], [0.0, 0.0]]) + if packed_verify: + logits = logits.repeat_interleave(torch.tensor([2, 4]), dim=0) + expected = expected.repeat_interleave(torch.tensor([2, 4]), dim=0) + result = SpecWorkerBase._zero_padding_rows( + logits, + metadata, + 1, + 3, + 1 if not packed_verify else 4, + packed_verify=packed_verify, + ) + assert torch.equal(result, expected) + + +def test_ragged_padding_rejects_incomplete_layout(): + metadata = _penalty_mapping_meta([3, 9], [2, 4]) + metadata.dummy_slot_row = 9 + metadata.total_verify_tokens = 5 + with pytest.raises(ValueError, match="verification layout"): + SpecWorkerBase._zero_padding_rows(torch.zeros(6, 8), metadata, 0, 2, 4, packed_verify=True) + + +@pytest.mark.parametrize("ragged", [False, True]) +def test_sampling_parameter_rows_follow_metadata_not_stale_request_windows(ragged): + config = types.SimpleNamespace(temperature=0.7, top_k=4, top_p=0.8, min_p=0.1) + requests = [ + types.SimpleNamespace( + sampling_config=config, state=state, py_seq_slot=slot, py_verify_len=window + ) + for slot, (state, window) in enumerate( + [ + (LlmRequestState.CONTEXT_INIT, 5), + (LlmRequestState.GENERATION_IN_PROGRESS, 1), + (LlmRequestState.GENERATION_IN_PROGRESS, 3), + ] + ) + ] + metadata = types.SimpleNamespace( + runtime_draft_len=5, + is_ragged_verify=ragged, + dummy_slot_row=9, + group_all_greedy_sample=None, + ) + normalized, slots = SpecMetadata._scan_one_model_sampling(metadata, requests) + assert [row[-1] for row in normalized] == ([1, 2, 4] if ragged else [1, 6, 6]) + assert [row[3] for row in normalized] == [0.1] * 3 + assert slots == [0, 1, 2] + + +@pytest.mark.parametrize( + "native, host, device", + [ + (True, True, False), + (True, False, True), + (False, True, True), + (1, False, False), + (False, 1, False), + ], +) +def test_conflicting_window_authorities_are_rejected(native, host, device): + with pytest.raises(ValueError): + SpecSampler._snapshot_policy_windows_for_step( + [], + native_uniform=native, + host_snapshot_required=host, + device_verify_lens_available=device, + ) + + +def test_native_uniform_marker_ignores_host_policy_window(): + request = types.SimpleNamespace(py_request_id=7, py_verify_len=2) + snapshot = SpecSampler._snapshot_policy_windows_for_step( + [request], + native_uniform=True, + host_snapshot_required=False, + device_verify_lens_available=False, + ) + request.py_verify_len = 1 + assert snapshot == {} + assert SpecSampler._verified_len(request, 5, snapshot) == 5 + + +def test_device_window_wins_over_host_shape_split(): + request = types.SimpleNamespace(py_request_id=7, py_verify_len=5) + snapshot = SpecSampler._snapshot_policy_windows_for_step( + [request], + native_uniform=False, + host_snapshot_required=False, + device_verify_lens_available=True, + ) + request.py_verify_len = 1 + assert snapshot is None + assert SpecSampler._verified_len(request, 5, snapshot, token_window=4) == 3 + + +def test_host_snapshot_includes_ragged_request_after_uniform_request(): + requests = [ + types.SimpleNamespace(py_request_id=7), + types.SimpleNamespace(py_request_id=8, py_verify_len=2), + ] + snapshot = SpecSampler._snapshot_policy_windows_for_step( + requests, + native_uniform=False, + host_snapshot_required=True, + device_verify_lens_available=False, + ) + requests[1].py_verify_len = 4 + assert snapshot == {8: 2} + assert SpecSampler._verified_len(requests[0], 5, snapshot) == 5 + assert SpecSampler._verified_len(requests[1], 5, snapshot) == 2 + + +@pytest.mark.parametrize("ragged", [False, True]) +def test_rejection_preserves_sampling_params_and_caps_acceptance(monkeypatch, ragged): + """Exercise packing around the backend boundary, without a CUDA kernel.""" + worker = _make_worker(0.0) + widths = [2, 4] if ragged else [4, 4] + metadata = _penalty_mapping_meta([4, 2, 9], widths if ragged else None) + metadata.dummy_slot_row = 9 + rows = 1 + sum(widths) + metadata.temperatures = torch.arange(rows, dtype=torch.float32) + 0.5 + metadata.top_ks = torch.arange(rows, dtype=torch.int32) + 1 + metadata.top_ps = torch.full((rows,), 0.9) + metadata.min_ps = torch.arange(rows, dtype=torch.float32) / 100 + metadata.advanced_sampling_mode = object() + logits = torch.arange(rows * 8, dtype=torch.float32).reshape(rows, 8) + logits[1 + widths[0] :] = float("nan") + calls = {} + + def compute_probs(mode, scores, temperatures, top_ks, top_ps, min_ps): + assert mode is metadata.advanced_sampling_mode + assert torch.equal(temperatures, metadata.temperatures[1:]) + assert torch.equal(top_ks, metadata.top_ks[1:]) + assert torch.equal(top_ps, metadata.top_ps[1:]) + assert torch.equal(min_ps, metadata.min_ps[1:]) + assert torch.equal(scores[: widths[0]], logits[1 : 1 + widths[0]]) + assert torch.equal(scores[widths[0] :], torch.zeros(widths[1], 8)) + return torch.softmax(scores, dim=-1) + + def rejection(**kwargs): + calls.update(kwargs) + return torch.zeros(2, 4, dtype=torch.int32), torch.tensor([4, 4]) + + worker._sample_tokens_for_batch = lambda *args: torch.tensor([7]) + worker._compute_probs_from_logits = compute_probs + worker._rng_state_per_request = lambda *args: (None, None) + monkeypatch.setattr(interface_ops, "rejection_sampling_one_model", rejection) + draft_probs = torch.full((2, 3, 8), 1 / 8) + original_probs = draft_probs.clone() + _, accepted = worker._sample_and_accept_draft_tokens_rejection( + logits, + torch.zeros(2, 3, dtype=torch.long), + draft_probs, + 1, + 3, + metadata, + ) + assert accepted.tolist() == [1, *widths] + assert torch.equal(draft_probs, original_probs) + if not ragged: + assert calls["draft_probs"] is draft_probs + target_probs = calls["target_probs"] + assert target_probs.shape == (2, 4, 8) + assert torch.allclose(target_probs.sum(-1), torch.ones(2, 4)) + if ragged: + assert target_probs[0, 2:, 0].tolist() == [1, 1] + assert target_probs[0, 2:, 1:].count_nonzero() == 0 + + +def _reference_rejection( + draft_probs: torch.Tensor, + draft_token_ids: torch.Tensor, + target_probs: torch.Tensor, + accept_uniforms: torch.Tensor, + recovery_uniforms: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """CPU probability-law oracle, not a native kernel or RNG implementation.""" + batch, width = draft_token_ids.shape + output = torch.full((batch, width + 1), -1, dtype=torch.int32) + lengths = torch.full((batch,), width + 1, dtype=torch.int32) + active = torch.ones(batch, dtype=torch.bool) + for pos in range(width + 1): + q = target_probs[:, pos] + if pos < width: + p = draft_probs[:, pos] + token = draft_token_ids[:, pos].long() + q_token = q.gather(1, token[:, None]).squeeze(1) + p_token = p.gather(1, token[:, None]).squeeze(1) + accept = active & (accept_uniforms * p_token < q_token) + output[:, pos] = torch.where(accept, token.int(), output[:, pos]) + residual = (q - p).clamp_min(0) + else: + accept = torch.zeros_like(active) + residual = q + cdf = (residual / residual.sum(-1, keepdim=True).clamp_min(1e-30)).cumsum(-1) + recovered = (recovery_uniforms[:, None] >= cdf).sum(-1).int() + rejected = active & ~accept + output[:, pos] = torch.where(rejected, recovered, output[:, pos]) + lengths = torch.where(rejected, pos + 1, lengths) + active = accept + return output, lengths + + +@pytest.mark.parametrize( + "windows, vocab_mode, preallocated, ragged", + [ + ([1], "full", False, True), + ([1, 2, 3, 4], "full", False, True), + ([1, 2, 3, 4], "prefix", False, True), + ([1, 2, 3, 4], "d2t", False, True), + ([1, 2, 3, 4], "d2t", True, True), + ([4, 4, 4, 4], "full", False, True), + ([4, 4, 4, 4], "full", False, False), + ], +) +def test_rejection_retains_request_local_bonus_distribution( + monkeypatch, windows, vocab_mode, preallocated, ragged +): + """All real drafts accept; the retained local bonus must still sample q.""" + width, vocab = 3, 4 + # Exact quadrature for these dyadic probabilities, not random sampling. + grid = (torch.arange(16, dtype=torch.float64) + 0.5) / 16 + trials, requests = grid.numel() ** 2, len(windows) + num_gens = trials * requests + accept_uniforms = grid.repeat_interleave(16 * requests) + recovery_uniforms = grid.repeat(16).repeat_interleave(requests) + bonuses = torch.tensor( + [[0.75, 0.25, 0, 0], [0.5, 0.25, 0.25, 0], [0, 0.25, 0.25, 0.5], [1, 0, 0, 0]] + )[:requests] + slots = torch.tensor([4, 1, 3, 2][:requests]).repeat(trials) + stored_vocab = vocab if vocab_mode == "full" else 2 + slot_probs = torch.zeros(6, width + 2, stored_vocab) + slot_probs[..., 0] = torch.arange(1, 7)[:, None] / 8 + slot_probs[..., 1] = 1 - slot_probs[..., 0] + expanded = torch.zeros(num_gens, width, vocab) + indices = torch.tensor([1, 3]) if vocab_mode == "d2t" else torch.arange(stored_vocab) + expanded[..., indices] = slot_probs[slots, :width] + worker = _make_worker() + if vocab_mode == "d2t": + # Neither draft-vocabulary ID maps to target token 0. + worker._d2t = torch.tensor([1, 2]) + worker._sample_tokens_for_batch = lambda *args: torch.tensor([3]) + worker._compute_probs_from_logits = lambda mode, scores, *params: scores.softmax(-1) + worker._rng_state_per_request = lambda *args: (None, None) + calls = {} + + def rejection(**kwargs): + calls.update(kwargs) + return _reference_rejection( + kwargs["draft_probs"], + kwargs["draft_token_ids"], + kwargs["target_probs"], + accept_uniforms, + recovery_uniforms, + ) + + monkeypatch.setattr(interface_ops, "rejection_sampling_one_model", rejection) + scratch = torch.zeros(num_gens + 2, width + 2, vocab) if preallocated else None + # Reuse expansion storage with different windows in the same token bucket. + for local_windows in [windows, windows[::-1]]: + lens = torch.tensor(local_windows * trials, dtype=torch.int32) + metadata = _penalty_mapping_meta([5, *slots.tolist()], lens.tolist() if ragged else None) + metadata.dummy_slot_row = 9 + metadata.runtime_draft_len = width + metadata.enable_penalty = False + metadata.use_rejection_sampling = True + metadata.is_all_greedy_sample = False + metadata.draft_probs = slot_probs + metadata.draft_probs_vocab_size = metadata.draft_probs_last_dim = stored_vocab + metadata.full_draft_probs = scratch + metadata.d2t_target_indices = None + packed = torch.cat( + [ + torch.cat([expanded[i, : length - 1], bonuses[i % requests, None]]) + for i, length in enumerate(lens.tolist()) + ] + ) + logits = torch.cat([torch.zeros(1, vocab), packed.log()]) + rows = logits.shape[0] + metadata.temperatures = torch.ones(rows) + metadata.top_ks = torch.full((rows,), vocab, dtype=torch.int32) + metadata.top_ps = torch.ones(rows) + metadata.min_ps = torch.zeros(rows) + metadata.advanced_sampling_mode = None + metadata.draft_tokens = torch.cat( + [torch.ones(int((lens - 1).sum()), dtype=torch.int32), torch.tensor([99, 99])] + ) + before = [t.clone() for t in (slot_probs, metadata.draft_tokens, logits, lens)] + + def no_host_read(*args, **kwargs): + raise AssertionError("device data must not be read by the host") + + with monkeypatch.context() as guard: + for name in ("item", "tolist", "__bool__", "numpy", "cpu"): + guard.setattr(torch.Tensor, name, no_host_read) + drafts = ( + interface_ops._padded_gen_draft_tokens(metadata, num_gens, width) + if ragged + else torch.ones(num_gens, width, dtype=torch.int32) + ) + drafts_before = drafts.clone() + output, lengths = worker._accept_draft_tokens(logits, drafts, 1, num_gens + 1, metadata) + + assert output[0, 0] == 3 and lengths[0] == 1 + assert torch.equal(lengths[1:], lens) + real = torch.arange(width)[None, :] < (lens - 1)[:, None] + assert torch.all(output[1:, :width][real] == 1) + bonus = output[1:].gather(1, (lens.long() - 1)[:, None]).reshape(trials, requests) + for request in range(requests): + observed = torch.bincount(bonus[:, request].long(), minlength=vocab) / trials + torch.testing.assert_close(observed, bonuses[request], rtol=0, atol=0) + passed_probs = calls["draft_probs"] + torch.testing.assert_close(passed_probs[real], expanded[real], rtol=0, atol=0) + if ragged: + expected_padding = torch.zeros_like(passed_probs[~real]) + expected_padding[:, 0] = 1 + torch.testing.assert_close(passed_probs[~real], expected_padding, rtol=0, atol=0) + else: + torch.testing.assert_close(passed_probs, expanded, rtol=0, atol=0) + for original, actual in zip(before, (slot_probs, metadata.draft_tokens, logits, lens)): + torch.testing.assert_close(actual, original, rtol=0, atol=0) + assert torch.equal(drafts, drafts_before) + assert torch.equal(calls["draft_token_ids"], drafts_before) + if scratch is not None: + # Synthetic mass at an unmapped ID must never survive in reusable storage. + torch.testing.assert_close(scratch[:num_gens, :width], expanded, rtol=0, atol=0) + assert not scratch[:, width:].count_nonzero() + assert not scratch[num_gens:].count_nonzero() + assert passed_probs.untyped_storage().data_ptr() != scratch.untyped_storage().data_ptr() + + +@pytest.fixture +def cpu_sampler(monkeypatch): + """Replace only CUDA transfers/events and the unrelated stop-criteria API.""" + original_to = torch.Tensor.to + + def to_cpu(tensor, *args, **kwargs): + if kwargs.get("device") == "cuda": + kwargs = {**kwargs, "device": "cpu"} + return original_to(tensor, *args, **kwargs) + + monkeypatch.setattr(torch.Tensor, "to", to_cpu) + monkeypatch.setattr(sampler_ops, "handle_stop_criteria", lambda *args, **kwargs: False) + sampler = SpecSampler.__new__(SpecSampler) + sampler.draft_len = 5 + sampler.max_accepted_path_len = 6 + sampler.max_seq_len = 128 + sampler.store = types.SimpleNamespace( + new_tokens=torch.zeros(6, 5, 1, dtype=torch.int32), + next_new_tokens=torch.zeros(6, 5, 1, dtype=torch.int32), + next_draft_tokens=torch.zeros(5, 5, dtype=torch.int32), + new_tokens_lens=torch.zeros(5, dtype=torch.int32), + ) + transfers = [] + + def copy_to_host(tensor): + transfers.append(("copy", tuple(tensor.shape))) + return tensor.clone() + + def record_event(): + transfers.append(("record", None)) + return types.SimpleNamespace(synchronize=lambda: transfers.append(("sync", None))) + + sampler._copy_to_host = copy_to_host + sampler._record_sampler_event = record_event + return sampler, transfers + + +def _sampling_step(): + def request(request_id, slot, context=False): + result = types.SimpleNamespace( + py_request_id=request_id, + py_seq_slot=slot, + state=LlmRequestState.CONTEXT_INIT + if context + else LlmRequestState.GENERATION_IN_PROGRESS, + py_draft_tokens=[] if context else [10, 11, 12], + py_draft_tokens_effective_len=None, + py_decoding_iter=0, + tokens=[], + ) + result.add_new_token = lambda token, beam_idx: result.tokens.append(token) + return result + + context, first, complete = request(7, 2, True), request(8, 3), request(9, 0) + first.py_verify_len = 2 + scheduled = types.SimpleNamespace( + context_requests_chunking=[object()], + context_requests_last_chunk=[context], + generation_requests=[first, complete], + ) + # Include both a skipped context chunk and an unsampled trailing graph row. + outputs = dict( + new_tokens=torch.arange(20, dtype=torch.int32).reshape(5, 4), + new_tokens_lens=torch.tensor([1, 1, 2, 2, 4], dtype=torch.int32), + next_draft_tokens=torch.arange(15, dtype=torch.int32).reshape(5, 3), + next_new_tokens=torch.arange(20, dtype=torch.int32).reshape(5, 4), + ) + return scheduled, outputs, context, first, complete + + +@pytest.mark.parametrize( + "authority, expected_verified, expected_rewind", + [("legacy", 2, 1), ("host", 2, 1), ("device", 1, 0), ("uniform", 3, 2)], +) +def test_step_windows_survive_request_and_output_reuse( + cpu_sampler, authority, expected_verified, expected_rewind +): + sampler, transfers = cpu_sampler + scheduled, outputs, context, first, complete = _sampling_step() + if authority == "device": + outputs["verify_lens"] = torch.tensor([99, 1, 2, 4, 99], dtype=torch.int32) + elif authority == "host": + outputs["host_policy_windows_snapshot"] = True + elif authority == "uniform": + outputs["native_uniform_verify"] = True + state = sampler.sample_async(scheduled, outputs, []) + assert state.requests == [context, first, complete] + assert state.draft_lens == [0, 3, 3] + assert state.runtime_draft_len == 3 + assert len(transfers) == (5 if authority == "device" else 4) + assert transfers[-1] == ("record", None) + if authority == "device": + assert state.host.verify_lens.tolist() == [1, 2, 4] + outputs["verify_lens"].fill_(6) + else: + assert state.host.verify_lens is None + # A following iteration mutates both the live request and slot-indexed store. + first.py_verify_len = 5 + first.py_draft_tokens = [99] * 5 + sampler.store.new_tokens.fill_(99) + complete.state = LlmRequestState.GENERATION_COMPLETE + sampler.update_requests(state) + assert transfers[-1] == ("sync", None) + assert first.tokens == [8, 9] + assert first.py_num_accepted_draft_tokens == 1 + assert first.py_num_draft_tokens_verified == expected_verified + assert first.py_rewind_len == expected_rewind + assert first.py_draft_tokens == [6, 7, 8] + assert first.py_decoding_iter == 1 + assert context.py_num_draft_tokens_verified == 0 + assert complete.tokens == [] + assert complete.py_decoding_iter == 0 + + +@pytest.mark.parametrize( + "invalid", + [ + [1, 2, 3, 4], + torch.ones(4, 1), + torch.ones(3, dtype=torch.int32), + torch.ones(4), + torch.ones(4, dtype=torch.bool), + ], +) +def test_device_windows_require_request_aligned_integer_vector(cpu_sampler, invalid): + sampler, transfers = cpu_sampler + scheduled, outputs, *_ = _sampling_step() + outputs["verify_lens"] = invalid + with pytest.raises(ValueError, match="verify_lens"): + sampler.sample_async(scheduled, outputs, []) + assert transfers == [] diff --git a/tests/unittest/_torch/speculative/test_capture_sampling_params.py b/tests/unittest/_torch/speculative/test_capture_sampling_params.py index b085424fb270..dd8d76071aed 100644 --- a/tests/unittest/_torch/speculative/test_capture_sampling_params.py +++ b/tests/unittest/_torch/speculative/test_capture_sampling_params.py @@ -108,7 +108,10 @@ def _request(temperature=None, top_k=None, top_p=None, min_p=None, slot=0): def _fake_meta(): return types.SimpleNamespace( - runtime_draft_len=1, dummy_slot_row=0, group_all_greedy_sample=None + runtime_draft_len=1, + is_ragged_verify=False, + dummy_slot_row=0, + group_all_greedy_sample=None, ) @@ -196,6 +199,7 @@ def _populate_meta(mode, draft_len=1): """ meta = types.SimpleNamespace( runtime_draft_len=draft_len, + is_ragged_verify=False, dummy_slot_row=0, group_all_greedy_sample=None, max_num_requests=4, diff --git a/tests/unittest/_torch/speculative/test_group_all_greedy_sync.py b/tests/unittest/_torch/speculative/test_group_all_greedy_sync.py index f679132618c0..8828060f5a73 100644 --- a/tests/unittest/_torch/speculative/test_group_all_greedy_sync.py +++ b/tests/unittest/_torch/speculative/test_group_all_greedy_sync.py @@ -36,6 +36,7 @@ def _fake_request(temperature=None, top_k=None, top_p=None, min_p=None, slot=0): def _fake_meta(group_all_greedy_sample=None): meta = types.SimpleNamespace( runtime_draft_len=2, + is_ragged_verify=False, dummy_slot_row=0, spec_dec_mode=types.SimpleNamespace(use_one_engine=lambda: True), group_all_greedy_sample=group_all_greedy_sample, diff --git a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py index 03151e690212..cdb7d0e1215c 100644 --- a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py +++ b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py @@ -90,6 +90,7 @@ def test_prepare_buffers_noop_when_disabled(): def _valid_state(): return types.SimpleNamespace( + is_ragged_verify=False, draft_probs=torch.empty((R, K, V), device="cuda"), batch_slot_ids=torch.arange(R, device="cuda", dtype=torch.long), ) @@ -208,6 +209,7 @@ def _forward_impl(self, *args: object, **kwargs: object) -> None: def _dispatch_meta(**over): base = dict( + is_ragged_verify=False, use_rejection_sampling=True, is_all_greedy_sample=False, draft_probs_vocab_size=V, From afc2dbd8f17ba7f14d7177a7f45cb2780fc740f1 Mon Sep 17 00:00:00 2001 From: Zhaoyuan He Date: Thu, 1 Oct 2026 18:07:13 -0700 Subject: [PATCH 2/3] [TRTLLM-14620][refactor] Use generic ragged helpers and typed metadata Access declared SpecMetadata fields directly and decouple the ragged helper module name from DSpark. Update callers and tests without changing helper arithmetic or production completion accounting. Add CPU regression cases for executed versus planned/accepted windows, pending-step storage ownership, and completed-request slot reuse. Preserve host snapshots and pre-forward planning until mixed row/context accounting can be migrated safely. Signed-off-by: Zhaoyuan He --- .../_torch/pyexecutor/sampler/penalties.py | 10 +-- tensorrt_llm/_torch/speculative/interface.py | 12 +-- .../{dspark_ragged.py => ragged_helpers.py} | 0 ...spark_ragged.py => test_ragged_helpers.py} | 2 +- ...ed_metadata.py => test_ragged_metadata.py} | 2 +- .../hw_agnostic/test_ragged_spec_sampler.py | 76 +++++++++++++++++++ 6 files changed, 89 insertions(+), 13 deletions(-) rename tensorrt_llm/_torch/speculative/{dspark_ragged.py => ragged_helpers.py} (100%) rename tests/unittest/_torch/speculative/hw_agnostic/{test_dspark_ragged.py => test_ragged_helpers.py} (99%) rename tests/unittest/_torch/speculative/hw_agnostic/{test_dspark_ragged_metadata.py => test_ragged_metadata.py} (97%) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/penalties.py b/tensorrt_llm/_torch/pyexecutor/sampler/penalties.py index bc498e486b22..244a26d12d83 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/penalties.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/penalties.py @@ -920,12 +920,12 @@ def build_row_mapping( num_gens = batch_size - num_contexts batch_slots = slot_ids[:batch_size].to(device=device, dtype=torch.int64) - if getattr(spec_metadata, "is_ragged_verify", False): - from tensorrt_llm._torch.speculative.dspark_ragged import row_ids_from_lens + if spec_metadata.is_ragged_verify: + from tensorrt_llm._torch.speculative.ragged_helpers import row_ids_from_lens - verify_lens = getattr(spec_metadata, "verify_lens", None) - qo_indptr = getattr(spec_metadata, "qo_indptr", None) - total_verify_tokens = getattr(spec_metadata, "total_verify_tokens", None) + verify_lens = spec_metadata.verify_lens + qo_indptr = spec_metadata.qo_indptr + total_verify_tokens = spec_metadata.total_verify_tokens if ( verify_lens is None or qo_indptr is None diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 6aceef382b55..8bfd073223c7 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -543,7 +543,7 @@ class SpecMetadata: # [qo_indptr[r], qo_indptr[r + 1]) # # Both stay None on every other path. Built by - # ``_torch/speculative/dspark_ragged.RaggedVerifyLayout``. + # ``_torch/speculative/ragged_helpers.RaggedVerifyLayout``. verify_lens: Optional[torch.Tensor] = None qo_indptr: Optional[torch.Tensor] = None # Host-side ``sum(verify_lens)`` (reading it off the tensor would sync). @@ -1592,7 +1592,7 @@ def _padded_gen_draft_tokens(spec_metadata: SpecMetadata, num_gens: int, rectangular, so scatter back out; positions past a request's window are padding and get masked by ``count_accepted_ragged``. """ - from .dspark_ragged import build_qo_indptr, scatter_ragged_to_padded + from .ragged_helpers import build_qo_indptr, scatter_ragged_to_padded # verify_lens counts tokens (bonus position + drafts); the draft buffer # only holds the drafts. @@ -2004,8 +2004,8 @@ def _sample_and_accept_draft_tokens_base( # Generation requests: verify draft tokens against target tokens if spec_metadata is not None and spec_metadata.is_ragged_verify: - from .dspark_ragged import (count_accepted_ragged, - scatter_ragged_to_padded) + from .ragged_helpers import (count_accepted_ragged, + scatter_ragged_to_padded) verify_lens = spec_metadata.verify_lens total = spec_metadata.total_verify_tokens qo_indptr = spec_metadata.qo_indptr @@ -2604,8 +2604,8 @@ def _sample_and_accept_draft_tokens_rejection( spec_metadata.top_ps[gen_start:gen_end], spec_metadata.min_ps[gen_start:gen_end]) if is_ragged: - from .dspark_ragged import (fill_padded_rows_onehot, - scatter_ragged_to_padded) + from .ragged_helpers import (fill_padded_rows_onehot, + scatter_ragged_to_padded) target_probs = scatter_ragged_to_padded( target_probs_flat, verify_lens=spec_metadata.verify_lens, diff --git a/tensorrt_llm/_torch/speculative/dspark_ragged.py b/tensorrt_llm/_torch/speculative/ragged_helpers.py similarity index 100% rename from tensorrt_llm/_torch/speculative/dspark_ragged.py rename to tensorrt_llm/_torch/speculative/ragged_helpers.py diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_helpers.py similarity index 99% rename from tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py rename to tests/unittest/_torch/speculative/hw_agnostic/test_ragged_helpers.py index 182cc5b0c221..65a905d947c0 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_helpers.py @@ -17,7 +17,7 @@ import pytest import torch -from tensorrt_llm._torch.speculative.dspark_ragged import ( +from tensorrt_llm._torch.speculative.ragged_helpers import ( RaggedVerifyLayout, build_qo_indptr, build_row_maps_device, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged_metadata.py b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_metadata.py similarity index 97% rename from tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged_metadata.py rename to tests/unittest/_torch/speculative/hw_agnostic/test_ragged_metadata.py index 89b294f36073..f655f00ed23c 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged_metadata.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_metadata.py @@ -18,7 +18,7 @@ import torch from tensorrt_llm._torch.pyexecutor.llm_request import get_request_tokens_per_gen_step -from tensorrt_llm._torch.speculative.dspark_ragged import ragged_gather_index_lists +from tensorrt_llm._torch.speculative.ragged_helpers import ragged_gather_index_lists pytestmark = pytest.mark.cpu_only diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py index 45fa4b430241..c3012158ab4a 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py @@ -691,6 +691,7 @@ def test_step_windows_survive_request_and_output_reuse( assert first.py_draft_tokens == [6, 7, 8] assert first.py_decoding_iter == 1 assert context.py_num_draft_tokens_verified == 0 + assert context.py_rewind_len == (0 if authority == "device" else 3) assert complete.tokens == [] assert complete.py_decoding_iter == 0 @@ -712,3 +713,78 @@ def test_device_windows_require_request_aligned_integer_vector(cpu_sampler, inva with pytest.raises(ValueError, match="verify_lens"): sampler.sample_async(scheduled, outputs, []) assert transfers == [] + + +@pytest.mark.parametrize( + "token_window, num_new_tokens", [(1, 1), (2, 1), (2, 2), (4, 1), (4, 2), (4, 4)] +) +def test_executed_device_window_is_not_host_plan_or_acceptance( + cpu_sampler, token_window, num_new_tokens +): + sampler, _ = cpu_sampler + scheduled, outputs, context, first, _ = _sampling_step() + # A device redistribution need not preserve the host's draft shape split. + first.py_verify_len = 0 + outputs["verify_lens"] = torch.tensor([99, 1, token_window, 4, 99], dtype=torch.int32) + outputs["new_tokens_lens"][2] = num_new_tokens + state = sampler.sample_async(scheduled, outputs, []) + first.py_verify_len = 5 + first.py_draft_tokens = [99] * 5 + sampler.update_requests(state) + + assert state.verify_lens_snapshot is None + assert first.py_num_accepted_draft_tokens == num_new_tokens - 1 + assert first.py_num_draft_tokens_verified == token_window - 1 + assert first.py_rewind_len == token_window - num_new_tokens + assert len(first.tokens) == num_new_tokens + assert context.py_num_draft_tokens_verified == 0 + assert context.py_rewind_len == 0 + + +def test_two_pending_steps_keep_separate_window_and_token_copies(cpu_sampler): + sampler, transfers = cpu_sampler + scheduled, outputs, _, first, _ = _sampling_step() + outputs["verify_lens"] = torch.tensor([99, 1, 4, 4, 99], dtype=torch.int32) + earlier = sampler.sample_async(scheduled, outputs, []) + + # The next step reuses both worker outputs and the slot-indexed sampler store. + outputs["verify_lens"][2] = 2 + outputs["new_tokens_lens"][2] = 1 + outputs["new_tokens"][2, 0] = 77 + later = sampler.sample_async(scheduled, outputs, []) + assert earlier.host.verify_lens.data_ptr() != later.host.verify_lens.data_ptr() + outputs["verify_lens"].fill_(99) + outputs["new_tokens"].fill_(99) + first.py_verify_len = 5 + + sampler.update_requests(earlier) + assert first.tokens == [8, 9] + assert first.py_num_draft_tokens_verified == 3 + assert first.py_rewind_len == 2 + sampler.update_requests(later) + assert first.tokens == [8, 9, 77] + assert first.py_num_accepted_draft_tokens == 0 + assert first.py_num_draft_tokens_verified == 1 + assert first.py_rewind_len == 1 + assert sum(event == "sync" for event, _ in transfers) == 2 + + +def test_completed_request_slot_reuse_does_not_rebind_pending_windows(cpu_sampler): + sampler, _ = cpu_sampler + scheduled, outputs, _, first, complete = _sampling_step() + outputs["verify_lens"] = torch.tensor([99, 1, 2, 4, 99], dtype=torch.int32) + earlier = sampler.sample_async(scheduled, outputs, []) + complete.state = LlmRequestState.GENERATION_COMPLETE + + next_scheduled, next_outputs, _, _, replacement = _sampling_step() + replacement.py_request_id = 10 + assert replacement.py_seq_slot == complete.py_seq_slot + next_outputs["verify_lens"] = torch.tensor([99, 1, 4, 2, 99], dtype=torch.int32) + later = sampler.sample_async(next_scheduled, next_outputs, []) + sampler.update_requests(earlier) + assert first.py_num_draft_tokens_verified == 1 + assert complete.tokens == [] + assert replacement.tokens == [] + sampler.update_requests(later) + assert replacement.py_num_draft_tokens_verified == 1 + assert replacement.py_rewind_len == 0 From 8109a50b2c4457c71936df3398362b2d97e5261f Mon Sep 17 00:00:00 2001 From: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> Date: Fri, 2 Oct 2026 05:13:18 -0700 Subject: [PATCH 3/3] [TRTLLM-14620][refactor] Consume executed output-row verification windows Publish context and generation verification lengths in persistent output-row-aligned device storage. Carry completed-step execution through the existing asynchronous per-step host copy instead of a request-ID Python snapshot. Preserve planning metadata, original proposal counts, native-uniform fallback and optional-attribute compatibility; reject incompatible legacy ragged outputs. Add ownership, alignment, acceptance and fallback regression tests. Signed-off-by: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/speculative/dspark.py | 65 ++- .../_torch/speculative/spec_sampler_base.py | 128 ++--- .../hw_agnostic/test_ragged_spec_sampler.py | 485 +++++++++++++++--- 3 files changed, 535 insertions(+), 143 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index b8dff7490d4c..203edc46a01b 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -53,6 +53,59 @@ def _dspark_position_ceiling(max_ctx: int, block_size: int, max_draft_len: int) return int(max_ctx) + int(max_draft_len) + int(block_size) + 3 +def _publish_policy_window_output( + outputs: dict[str, torch.Tensor | bool], + verify_lens: Optional[torch.Tensor], + *, + num_contexts: int, + num_generations: int, + output_verify_lens: Optional[torch.Tensor], +) -> None: + """Publish executed token windows in full output-row order. + + ``verify_lens`` is generation-ordered, including graph-padding generations. + ``output_verify_lens`` is persistent int32 device storage sized at warmup. + Context rows have one anchor and no verified drafts. The sampler removes + skipped context rows and snapshots the remaining view through its D2H path. + """ + batch_size = num_contexts + num_generations + if num_contexts < 0 or num_generations < 0: + raise ValueError("verification output row counts must be nonnegative") + for name in ("new_tokens", "new_tokens_lens", "next_draft_tokens", "next_new_tokens"): + if outputs[name].shape[0] != batch_size: + raise ValueError("verification output row counts must match every worker output") + + if verify_lens is None: + outputs.pop("verify_lens", None) + outputs.pop("verify_lens_in_output_order", None) + outputs.pop("host_policy_windows_snapshot", None) + outputs["native_uniform_verify"] = True + return + if ( + verify_lens.dim() != 1 + or verify_lens.shape[0] != num_generations + or verify_lens.dtype != torch.int32 + or verify_lens.device != outputs["new_tokens_lens"].device + ): + raise ValueError("verify_lens must contain exactly one int32 window per generation row") + if ( + output_verify_lens is None + or output_verify_lens.dim() != 1 + or output_verify_lens.shape[0] < batch_size + or output_verify_lens.dtype != torch.int32 + or output_verify_lens.device != verify_lens.device + ): + raise ValueError("executed-window output storage must cover the full batch on device") + + output_view = output_verify_lens[:batch_size] + output_view[:num_contexts].fill_(1) + output_view[num_contexts:].copy_(verify_lens) + outputs.pop("native_uniform_verify", None) + outputs.pop("host_policy_windows_snapshot", None) + outputs["verify_lens"] = output_view + outputs["verify_lens_in_output_order"] = True + + @dataclass class DSparkSpecMetadata(SpecMetadata): """Metadata for DSpark speculative decoding. @@ -267,6 +320,7 @@ def __init__( # Per-slot rolling captured-context KV windows, built lazily on the # first forward (fixed-size for slot-indexed reads/writes). self._win_inited = False + self._output_verify_lens: Optional[torch.Tensor] = None self._attention_warmup_attempted = False self._kv_windows: Optional[torch.Tensor] = None # [max_batch, num_stages, win, hd] self._ctx_len: Optional[torch.Tensor] = None # [max_batch] abs decode position @@ -354,6 +408,7 @@ def _lazy_init(self, draft_model, spec_metadata, attn_metadata=None) -> None: if not self._win_inited: max_batch = spec_metadata.max_num_requests + self._output_verify_lens = torch.empty(max_batch, dtype=torch.int32, device="cuda") num_stages = draft_model.num_stages self._win = int(draft_model._attn_params["window_size"]) head_dim = int(draft_model._attn_params["head_dim"]) @@ -819,13 +874,21 @@ def _forward_impl( self._position_initialized.copy_(saved_position_initialized) self._kv_windows.copy_(saved_windows) - return { + outputs = { "logits": raw_logits, "new_tokens": accepted_tokens, "new_tokens_lens": num_accepted_tokens, "next_draft_tokens": next_draft_tokens, "next_new_tokens": next_new_tokens, } + _publish_policy_window_output( + outputs, + spec_metadata.verify_lens, + num_contexts=num_contexts, + num_generations=num_gens, + output_verify_lens=self._output_verify_lens, + ) + return outputs class DSparkWorker(DFlashWorker): diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 92f1305c1af2..5fb5b185dc0b 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -148,10 +148,6 @@ class SampleStateSpec(SampleState): # instead would see the NEXT step's buffer, which update_requests itself # installs. draft_lens: Optional[list[int]] = None - # Per-request verify windows for this sampler step, keyed by request id. - # The overlap scheduler can stamp the next step on the live request before - # this state is consumed, so rewind accounting must use this snapshot. - verify_lens_snapshot: Optional[dict[int, int]] = None class SpecSampler(Sampler[SampleStateSpec], AsyncWorkerMixin): @@ -327,57 +323,6 @@ def _request_common_handling( request.py_draft_tokens = next_draft_tokens[request.py_seq_slot][:runtime_draft_len] request.py_decoding_iter += 1 - @staticmethod - def _verified_len( - request: LlmRequest, - runtime_draft_len: int, - verify_lens_snapshot: Optional[dict[int, int]], - token_window: Optional[int] = None, - ) -> int: - """Return the number of draft positions verified for one request.""" - if token_window is not None and token_window > 0: - return int(token_window) - 1 - verify_len = ( - verify_lens_snapshot.get(request.py_request_id) - if verify_lens_snapshot is not None - else None - ) - return runtime_draft_len if verify_len is None else int(verify_len) - - @staticmethod - def _snapshot_verify_lens( - requests: list[LlmRequest], - ) -> Optional[dict[int, int]]: - """Capture overlap-sensitive ragged windows for one sampler step.""" - snapshot = { - request.py_request_id: int(verify_len) - for request in requests - if (verify_len := getattr(request, "py_verify_len", None)) is not None - } - return snapshot or None - - @classmethod - def _snapshot_policy_windows_for_step( - cls, - requests: list[LlmRequest], - *, - native_uniform: bool, - host_snapshot_required: bool, - device_verify_lens_available: bool, - ) -> Optional[dict[int, int]]: - """Select one authoritative verification-window source for this step.""" - if type(native_uniform) is not bool or type(host_snapshot_required) is not bool: - raise ValueError("verification-window markers must be boolean") - if native_uniform and (host_snapshot_required or device_verify_lens_available): - raise ValueError("native-uniform verification cannot also publish ragged windows") - if host_snapshot_required and device_verify_lens_available: - raise ValueError("device windows and host-window snapshots are mutually exclusive") - if native_uniform: - return {} - if device_verify_lens_available: - return None - return cls._snapshot_verify_lens(requests) - def update_requests( self, state: SampleStateSpec, @@ -403,6 +348,23 @@ def update_requests( beam_idx = DEFAULT_BEAM_IDX runtime_draft_len = getattr(state, "runtime_draft_len", self.draft_len) + if token_windows is not None: + for request_index, request in enumerate(state.requests): + if request.state == LlmRequestState.GENERATION_COMPLETE: + continue + token_window = token_windows[request_index] + num_new_tokens = new_tokens_lens_list[request.py_seq_slot] + max_new_tokens = min(token_window, self.max_accepted_path_len) + if state.draft_lens is not None: + max_new_tokens = min(max_new_tokens, state.draft_lens[request_index] + 1) + if not 1 <= token_window <= runtime_draft_len + 1 or not ( + 1 <= num_new_tokens <= max_new_tokens + ): + raise ValueError( + "executed verify_lens must be positive and new_tokens_lens must fit " + "the executed window, original proposals, and token storage" + ) + for req_idx, req in enumerate(state.requests): if req.state == LlmRequestState.GENERATION_COMPLETE: continue @@ -427,14 +389,7 @@ def update_requests( # length must not be used as the denominator (0 for the request's # prefill step, where nothing was verified). verified_len = ( - runtime_draft_len - if token_windows is None and state.verify_lens_snapshot is None - else self._verified_len( - req, - runtime_draft_len, - state.verify_lens_snapshot, - token_windows[req_idx] if token_windows is not None else None, - ) + token_windows[req_idx] - 1 if token_windows is not None else runtime_draft_len ) drafted_len = state.draft_lens[req_idx] if state.draft_lens is not None else 0 req.py_num_draft_tokens_verified = min(drafted_len, verified_len) @@ -458,9 +413,10 @@ def sample_async( - new_tokens_lens: [batch] number of accepted tokens - next_draft_tokens: [batch, max_draft_len] draft tokens for next iter - next_new_tokens: [batch, max_draft_len + 1] input for next iter - - verify_lens: optional [batch] device-selected token windows + - verify_lens: optional int32 [output rows] executed token windows + - verify_lens_in_output_order: required True with verify_lens - native_uniform_verify: optional boolean declaring uniform windows - - host_policy_windows_snapshot: optional boolean declaring host windows + - host_policy_windows_snapshot: legacy True is rejected num_context_logits_prefix_sum: Prefix sum of context logits (unused) Returns: @@ -498,23 +454,40 @@ def sample_async( runtime_draft_len = o_next_draft_tokens.shape[1] o_verify_lens = outputs.get("verify_lens") - if o_verify_lens is not None: + native_uniform = outputs.get("native_uniform_verify", False) + host_snapshot = outputs.get("host_policy_windows_snapshot", False) + output_order = outputs.get("verify_lens_in_output_order", False) + if any( + type(marker) is not bool for marker in (native_uniform, host_snapshot, output_order) + ): + raise ValueError("verification-window markers must be boolean") + if host_snapshot: + raise ValueError("legacy ragged producers must publish executed output-row verify_lens") + if native_uniform and (o_verify_lens is not None or output_order): + raise ValueError("native-uniform verification cannot also publish ragged windows") + if output_order != (o_verify_lens is not None): + raise ValueError("verify_lens requires an explicit output-row layout declaration") + if o_verify_lens is None: + if not native_uniform and any( + getattr(r, "py_verify_len", None) is not None for r in sampling_requests + ): + raise ValueError("ragged producers must publish executed output-row verify_lens") + else: + output_rows = outputs["new_tokens_lens"].shape[0] if ( not isinstance(o_verify_lens, torch.Tensor) or o_verify_lens.dim() != 1 - or o_verify_lens.dtype not in (torch.int32, torch.int64) - or o_verify_lens.shape[0] < num_skip + num_sampling_requests - ): - raise ValueError( - "verify_lens must contain one integer token window per sampled request" + or o_verify_lens.dtype != torch.int32 + or o_verify_lens.device != outputs["new_tokens_lens"].device + or o_verify_lens.shape[0] != output_rows + or output_rows < num_skip + num_sampling_requests + or any( + outputs[name].shape[0] != output_rows + for name in ("new_tokens", "next_draft_tokens", "next_new_tokens") ) + ): + raise ValueError("verify_lens must contain one int32 token window per output row") o_verify_lens = o_verify_lens[num_skip : num_skip + num_sampling_requests] - verify_lens_snapshot = self._snapshot_policy_windows_for_step( - sampling_requests, - native_uniform=outputs.get("native_uniform_verify", False), - host_snapshot_required=outputs.get("host_policy_windows_snapshot", False), - device_verify_lens_available=o_verify_lens is not None, - ) # Pad or truncate to match fixed-size store buffers for index_copy_. # The worker output width tracks runtime_draft_len, which dynamic draft @@ -573,5 +546,4 @@ def sample_async( sampler_event=sampler_event, runtime_draft_len=runtime_draft_len, draft_lens=draft_lens, - verify_lens_snapshot=verify_lens_snapshot, ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py index c3012158ab4a..7afbea6046a3 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py @@ -166,22 +166,22 @@ def _apply(logits, spec_metadata, row_slots, intra_tokens, intra_valid): assert captured["intra_valid"].sum().item() == 7 -def test_verify_window_snapshot_wins_over_next_overlap_step(): - request = types.SimpleNamespace(py_request_id=7, py_verify_len=2) - snapshot = SpecSampler._snapshot_verify_lens([request]) - request.py_verify_len = 5 - - assert snapshot == {7: 2} - assert SpecSampler._verified_len(request, 5, snapshot) == 2 - +def test_legacy_ragged_without_executed_windows_is_rejected(cpu_sampler): + sampler, _ = cpu_sampler + scheduled, outputs, *_ = _sampling_step() + with pytest.raises(ValueError, match="ragged producers"): + sampler.sample_async(scheduled, outputs, []) -def test_uniform_verify_window_keeps_runtime_draft_length(): - request = types.SimpleNamespace(py_request_id=7) - snapshot = SpecSampler._snapshot_verify_lens([request]) - request.py_verify_len = 1 - assert snapshot is None - assert SpecSampler._verified_len(request, 5, snapshot) == 5 +def test_uniform_verify_window_keeps_runtime_draft_length(cpu_sampler): + sampler, _ = cpu_sampler + scheduled, outputs, _, first, _ = _sampling_step() + outputs["native_uniform_verify"] = True + state = sampler.sample_async(scheduled, outputs, []) + sampler.update_requests(state) + assert state.host.verify_lens is None + assert first.py_num_draft_tokens_verified == 3 + assert first.py_rewind_len == 2 @pytest.mark.parametrize("force_accept", [0.0, 1.0, 3.0]) @@ -309,57 +309,50 @@ def test_sampling_parameter_rows_follow_metadata_not_stale_request_windows(ragge (False, 1, False), ], ) -def test_conflicting_window_authorities_are_rejected(native, host, device): +def test_conflicting_window_authorities_are_rejected(cpu_sampler, native, host, device): + sampler, _ = cpu_sampler + scheduled, outputs, *_ = _sampling_step() + outputs["native_uniform_verify"] = native + outputs["host_policy_windows_snapshot"] = host + if device: + outputs["verify_lens"] = torch.tensor([1, 1, 3, 4, 1], dtype=torch.int32) + outputs["verify_lens_in_output_order"] = True with pytest.raises(ValueError): - SpecSampler._snapshot_policy_windows_for_step( - [], - native_uniform=native, - host_snapshot_required=host, - device_verify_lens_available=device, - ) + sampler.sample_async(scheduled, outputs, []) -def test_native_uniform_marker_ignores_host_policy_window(): - request = types.SimpleNamespace(py_request_id=7, py_verify_len=2) - snapshot = SpecSampler._snapshot_policy_windows_for_step( - [request], - native_uniform=True, - host_snapshot_required=False, - device_verify_lens_available=False, - ) - request.py_verify_len = 1 - assert snapshot == {} - assert SpecSampler._verified_len(request, 5, snapshot) == 5 - - -def test_device_window_wins_over_host_shape_split(): - request = types.SimpleNamespace(py_request_id=7, py_verify_len=5) - snapshot = SpecSampler._snapshot_policy_windows_for_step( - [request], - native_uniform=False, - host_snapshot_required=False, - device_verify_lens_available=True, - ) - request.py_verify_len = 1 - assert snapshot is None - assert SpecSampler._verified_len(request, 5, snapshot, token_window=4) == 3 +def test_native_uniform_marker_ignores_host_policy_window(cpu_sampler): + sampler, _ = cpu_sampler + scheduled, outputs, _, first, _ = _sampling_step() + outputs["native_uniform_verify"] = True + state = sampler.sample_async(scheduled, outputs, []) + first.py_verify_len = 1 + sampler.update_requests(state) + assert state.host.verify_lens is None + assert first.py_num_draft_tokens_verified == 3 -def test_host_snapshot_includes_ragged_request_after_uniform_request(): - requests = [ - types.SimpleNamespace(py_request_id=7), - types.SimpleNamespace(py_request_id=8, py_verify_len=2), - ] - snapshot = SpecSampler._snapshot_policy_windows_for_step( - requests, - native_uniform=False, - host_snapshot_required=True, - device_verify_lens_available=False, - ) - requests[1].py_verify_len = 4 - assert snapshot == {8: 2} - assert SpecSampler._verified_len(requests[0], 5, snapshot) == 5 - assert SpecSampler._verified_len(requests[1], 5, snapshot) == 2 +def test_device_window_wins_over_host_shape_split(cpu_sampler): + sampler, _ = cpu_sampler + scheduled, outputs, _, first, _ = _sampling_step() + first.py_verify_len = 5 + outputs["verify_lens"] = torch.tensor([1, 1, 4, 3, 1], dtype=torch.int32) + outputs["verify_lens_in_output_order"] = True + state = sampler.sample_async(scheduled, outputs, []) + first.py_verify_len = 1 + sampler.update_requests(state) + assert first.py_num_draft_tokens_verified == 3 + assert first.py_rewind_len == 2 + + +def test_legacy_host_snapshot_marker_is_rejected_before_request_updates(cpu_sampler): + sampler, transfers = cpu_sampler + scheduled, outputs, context, first, last = _sampling_step() + outputs["host_policy_windows_snapshot"] = True + with pytest.raises(ValueError, match="legacy ragged producers"): + sampler.sample_async(scheduled, outputs, []) + assert transfers == [] + assert context.tokens == first.tokens == last.tokens == [] @pytest.mark.parametrize("ragged", [False, True]) @@ -628,6 +621,7 @@ def request(request_id, slot, context=False): else LlmRequestState.GENERATION_IN_PROGRESS, py_draft_tokens=[] if context else [10, 11, 12], py_draft_tokens_effective_len=None, + py_verify_len=None, py_decoding_iter=0, tokens=[], ) @@ -653,7 +647,7 @@ def request(request_id, slot, context=False): @pytest.mark.parametrize( "authority, expected_verified, expected_rewind", - [("legacy", 2, 1), ("host", 2, 1), ("device", 1, 0), ("uniform", 3, 2)], + [("device", 1, 0), ("uniform", 3, 2)], ) def test_step_windows_survive_request_and_output_reuse( cpu_sampler, authority, expected_verified, expected_rewind @@ -662,8 +656,7 @@ def test_step_windows_survive_request_and_output_reuse( scheduled, outputs, context, first, complete = _sampling_step() if authority == "device": outputs["verify_lens"] = torch.tensor([99, 1, 2, 4, 99], dtype=torch.int32) - elif authority == "host": - outputs["host_policy_windows_snapshot"] = True + outputs["verify_lens_in_output_order"] = True elif authority == "uniform": outputs["native_uniform_verify"] = True state = sampler.sample_async(scheduled, outputs, []) @@ -710,6 +703,7 @@ def test_device_windows_require_request_aligned_integer_vector(cpu_sampler, inva sampler, transfers = cpu_sampler scheduled, outputs, *_ = _sampling_step() outputs["verify_lens"] = invalid + outputs["verify_lens_in_output_order"] = True with pytest.raises(ValueError, match="verify_lens"): sampler.sample_async(scheduled, outputs, []) assert transfers == [] @@ -726,13 +720,14 @@ def test_executed_device_window_is_not_host_plan_or_acceptance( # A device redistribution need not preserve the host's draft shape split. first.py_verify_len = 0 outputs["verify_lens"] = torch.tensor([99, 1, token_window, 4, 99], dtype=torch.int32) + outputs["verify_lens_in_output_order"] = True outputs["new_tokens_lens"][2] = num_new_tokens state = sampler.sample_async(scheduled, outputs, []) first.py_verify_len = 5 first.py_draft_tokens = [99] * 5 sampler.update_requests(state) - assert state.verify_lens_snapshot is None + assert not hasattr(state, "verify_lens_snapshot") assert first.py_num_accepted_draft_tokens == num_new_tokens - 1 assert first.py_num_draft_tokens_verified == token_window - 1 assert first.py_rewind_len == token_window - num_new_tokens @@ -745,6 +740,7 @@ def test_two_pending_steps_keep_separate_window_and_token_copies(cpu_sampler): sampler, transfers = cpu_sampler scheduled, outputs, _, first, _ = _sampling_step() outputs["verify_lens"] = torch.tensor([99, 1, 4, 4, 99], dtype=torch.int32) + outputs["verify_lens_in_output_order"] = True earlier = sampler.sample_async(scheduled, outputs, []) # The next step reuses both worker outputs and the slot-indexed sampler store. @@ -773,6 +769,7 @@ def test_completed_request_slot_reuse_does_not_rebind_pending_windows(cpu_sample sampler, _ = cpu_sampler scheduled, outputs, _, first, complete = _sampling_step() outputs["verify_lens"] = torch.tensor([99, 1, 2, 4, 99], dtype=torch.int32) + outputs["verify_lens_in_output_order"] = True earlier = sampler.sample_async(scheduled, outputs, []) complete.state = LlmRequestState.GENERATION_COMPLETE @@ -780,6 +777,7 @@ def test_completed_request_slot_reuse_does_not_rebind_pending_windows(cpu_sample replacement.py_request_id = 10 assert replacement.py_seq_slot == complete.py_seq_slot next_outputs["verify_lens"] = torch.tensor([99, 1, 4, 2, 99], dtype=torch.int32) + next_outputs["verify_lens_in_output_order"] = True later = sampler.sample_async(next_scheduled, next_outputs, []) sampler.update_requests(earlier) assert first.py_num_draft_tokens_verified == 1 @@ -788,3 +786,362 @@ def test_completed_request_slot_reuse_does_not_rebind_pending_windows(cpu_sample sampler.update_requests(later) assert replacement.py_num_draft_tokens_verified == 1 assert replacement.py_rewind_len == 0 + + +@pytest.mark.parametrize( + "output_row, token_window, num_new_tokens, effective_len", + [ + (2, 0, 2, None), + (2, -1, 2, None), + (2, 1, 2, None), + (2, 2, 0, None), + (2, 2, -1, None), + (3, 1, 2, None), + (2, 4, 3, 1), + (1, 4, 2, None), + (3, 8, 7, None), + ], +) +def test_invalid_executed_windows_fail_before_request_updates( + cpu_sampler, output_row, token_window, num_new_tokens, effective_len +): + sampler, _ = cpu_sampler + scheduled, outputs, context, first, last = _sampling_step() + first.py_draft_tokens_effective_len = effective_len + outputs["verify_lens"] = torch.tensor([99, 1, 3, 4, 99], dtype=torch.int32) + outputs["verify_lens_in_output_order"] = True + outputs["verify_lens"][output_row] = token_window + outputs["new_tokens_lens"][output_row] = num_new_tokens + state = sampler.sample_async(scheduled, outputs, []) + + with pytest.raises(ValueError, match="executed verify_lens"): + sampler.update_requests(state) + + assert context.tokens == first.tokens == last.tokens == [] + assert context.py_decoding_iter == first.py_decoding_iter == last.py_decoding_iter == 0 + + +def test_invalid_executed_window_for_completed_request_is_not_consumed(cpu_sampler): + sampler, _ = cpu_sampler + scheduled, outputs, _, first, complete = _sampling_step() + outputs["verify_lens"] = torch.tensor([99, 1, 3, -1, 99], dtype=torch.int32) + outputs["verify_lens_in_output_order"] = True + state = sampler.sample_async(scheduled, outputs, []) + complete.state = LlmRequestState.GENERATION_COMPLETE + sampler.update_requests(state) + + assert first.py_num_draft_tokens_verified == 2 + assert first.py_rewind_len == 1 + assert complete.tokens == [] + + +@pytest.mark.parametrize("contexts,generations", [(0, 3), (2, 3), (3, 0), (0, 0)]) +def test_publisher_uses_explicit_row_counts_without_host_sync(monkeypatch, contexts, generations): + from tensorrt_llm._torch.speculative.dspark import _publish_policy_window_output + + batch = contexts + generations + outputs = { + name: torch.zeros((batch, 4), dtype=torch.int32) + for name in ("new_tokens", "next_new_tokens", "next_draft_tokens") + } + outputs["new_tokens_lens"] = torch.ones(batch, dtype=torch.int32) + executed = torch.arange(1, generations + 1, dtype=torch.int32) + storage = torch.full((batch + 3,), -1, dtype=torch.int32) + pointer = storage.data_ptr() + + def forbidden(*args, **kwargs): + raise AssertionError("publisher must not allocate or read device values on host") + + with monkeypatch.context() as guard: + for name in ("item", "tolist", "__bool__", "numpy", "cpu", "clone", "to"): + guard.setattr(torch.Tensor, name, forbidden) + for name in ("empty", "zeros", "ones", "full", "tensor", "cat", "stack"): + guard.setattr(torch, name, forbidden) + _publish_policy_window_output( + outputs, + executed, + num_contexts=contexts, + num_generations=generations, + output_verify_lens=storage, + ) + assert outputs["verify_lens_in_output_order"] is True + assert outputs["verify_lens"].tolist() == [1] * contexts + list(range(1, generations + 1)) + assert storage.data_ptr() == pointer + assert storage[batch:].tolist() == [-1] * 3 + if batch: + assert outputs["verify_lens"].data_ptr() == pointer + + +@pytest.mark.parametrize( + "invalid", ["capacity", "short", "dtype", "dimensions", "storage", "missing_storage", "rows"] +) +def test_publisher_rejects_ambiguous_capacity_and_bad_shapes(invalid): + from tensorrt_llm._torch.speculative.dspark import _publish_policy_window_output + + _, outputs, *_ = _sampling_step() + executed = torch.tensor([3, 4, 1], dtype=torch.int32) + storage = torch.empty(8, dtype=torch.int32) + if invalid == "capacity": + executed = torch.tensor([3, 4, 1, 99, 99], dtype=torch.int32) + elif invalid == "short": + executed = executed[:2] + elif invalid == "dtype": + executed = executed.long() + elif invalid == "dimensions": + executed = executed[:, None] + elif invalid == "storage": + storage = storage[:4] + elif invalid == "missing_storage": + storage = None + else: + outputs["next_new_tokens"] = outputs["next_new_tokens"][:4] + with pytest.raises(ValueError): + _publish_policy_window_output( + outputs, executed, num_contexts=2, num_generations=3, output_verify_lens=storage + ) + + +def test_publisher_consumer_mixed_skipped_sparse_padding_and_stale_plan(cpu_sampler): + from tensorrt_llm._torch.speculative.dspark import _publish_policy_window_output + + sampler, _ = cpu_sampler + scheduled, outputs, context, first, second = _sampling_step() + first.py_verify_len, second.py_verify_len = 0, 1 + storage = torch.empty(8, dtype=torch.int32) + _publish_policy_window_output( + outputs, + torch.tensor([3, 4, 1], dtype=torch.int32), + num_contexts=2, + num_generations=3, + output_verify_lens=storage, + ) + assert outputs["verify_lens"].tolist() == [1, 1, 3, 4, 1] + state = sampler.sample_async(scheduled, outputs, []) + assert [r.py_seq_slot for r in state.requests] == [2, 3, 0] + assert state.host.verify_lens.tolist() == [1, 3, 4] + assert state.draft_lens == [0, 3, 3] + first.py_verify_len, second.py_verify_len = 5, 0 + storage.fill_(99) + sampler.update_requests(state) + assert [r.py_rewind_len for r in state.requests] == [0, 1, 2] + assert [r.py_num_draft_tokens_verified for r in state.requests] == [0, 2, 3] + assert [r.py_num_accepted_draft_tokens for r in state.requests] == [0, 1, 1] + assert [r.tokens for r in state.requests] == [[4], [8, 9], [12, 13]] + + +def test_publisher_two_pending_steps_reuse_storage_not_host_ownership(cpu_sampler): + from tensorrt_llm._torch.speculative.dspark import _publish_policy_window_output + + sampler, _ = cpu_sampler + storage = torch.empty(8, dtype=torch.int32) + states = [] + for generation_windows in ([3, 4, 1], [4, 2, 1]): + scheduled, outputs, *_ = _sampling_step() + _publish_policy_window_output( + outputs, + torch.tensor(generation_windows, dtype=torch.int32), + num_contexts=2, + num_generations=3, + output_verify_lens=storage, + ) + states.append(sampler.sample_async(scheduled, outputs, [])) + storage.fill_(99) + assert states[0].host.verify_lens.tolist() == [1, 3, 4] + assert states[1].host.verify_lens.tolist() == [1, 4, 2] + assert states[0].host.verify_lens.data_ptr() != states[1].host.verify_lens.data_ptr() + for state in states: + sampler.update_requests(state) + assert [r.py_rewind_len for r in states[0].requests] == [0, 1, 2] + assert [r.py_rewind_len for r in states[1].requests] == [0, 2, 0] + + +@pytest.mark.parametrize( + "stale_marker", ["host_policy_windows_snapshot", "verify_lens_in_output_order"] +) +def test_publisher_native_fallback_clears_stale_ragged_authority(cpu_sampler, stale_marker): + from tensorrt_llm._torch.speculative.dspark import _publish_policy_window_output + + sampler, _ = cpu_sampler + scheduled, outputs, _, first, _ = _sampling_step() + outputs[stale_marker] = True + outputs["verify_lens"] = torch.full((5,), 99, dtype=torch.int32) + _publish_policy_window_output( + outputs, None, num_contexts=2, num_generations=3, output_verify_lens=None + ) + assert outputs["native_uniform_verify"] is True + assert "verify_lens" not in outputs and "verify_lens_in_output_order" not in outputs + assert "host_policy_windows_snapshot" not in outputs + state = sampler.sample_async(scheduled, outputs, []) + assert state.host.verify_lens is None + sampler.update_requests(state) + assert first.py_rewind_len == 2 + assert first.py_num_draft_tokens_verified == 3 + + +@pytest.mark.parametrize("return_confidence", [False, True]) +@pytest.mark.parametrize("executed_window", [None, 2]) +def test_dspark_forward_publishes_execution_not_confidence_flag(return_confidence, executed_window): + from tensorrt_llm._torch.speculative.dspark import DSv4DSparkWorker + + worker = DSv4DSparkWorker.__new__(DSv4DSparkWorker) + worker.spec_config = types.SimpleNamespace(max_draft_len=1) + worker.return_confidence = return_confidence + worker.guided_decoder = None + worker._output_verify_lens = torch.empty(4, dtype=torch.int32) + worker._lazy_init = lambda *_args: None + worker._execute_guided_decoder_if_present = lambda *_args: None + worker.sample_and_accept_draft_tokens = lambda *_args: ( + torch.tensor([[11, 12]]), + torch.tensor([2]), + ) + worker._draft_gen_block_batched = lambda *_args, **_kwargs: torch.zeros((1, 1, 8)) + worker.sample_draft_tokens = lambda *_args, **_kwargs: torch.tensor([[13]]) + worker.write_context_onehot_draft_probs = lambda *_args: None + worker._prepare_next_new_tokens = lambda *_args: torch.tensor([[12, 13]]) + metadata = types.SimpleNamespace( + batch_indices_cuda=torch.tensor([0]), + draft_probs_last_dim=8, + verify_lens=None + if executed_window is None + else torch.tensor([executed_window], dtype=torch.int32), + ) + + outputs = worker._forward_impl( + input_ids=torch.tensor([11, 12]), + position_ids=torch.tensor([[0, 1]]), + hidden_states=None, + logits=torch.zeros((2, 8)), + attn_metadata=types.SimpleNamespace(num_seqs=1, num_contexts=0), + spec_metadata=metadata, + draft_model=types.SimpleNamespace(block_size=1), + ) + + if executed_window is None: + assert outputs["native_uniform_verify"] is True + assert "verify_lens" not in outputs + else: + assert outputs["verify_lens"].tolist() == [executed_window] + assert outputs["verify_lens_in_output_order"] is True + assert "native_uniform_verify" not in outputs + + +@pytest.mark.parametrize("deferred", [False, True]) +def test_publisher_exact_copy_owns_two_pending_steps(cpu_sampler, monkeypatch, deferred): + from concurrent.futures import Future + + from tensorrt_llm._torch.pyexecutor.sampler.sampler_features import AsyncWorkerMixin + from tensorrt_llm._torch.speculative.dspark import _publish_policy_window_output + + class Event: + def record(self): + self.recorded = True + + def synchronize(self): + assert self.recorded + + class DeferredCopies: + def __init__(self): + self.pending = [] + + def submit(self, function, *args): + future = Future() + self.pending.append((future, function, args)) + return future + + def finish(self): + for future, function, args in self.pending: + future.set_result(function(*args)) + + monkeypatch.setattr(torch.cuda, "Event", Event) + original_empty_like = torch.empty_like + monkeypatch.setattr( + torch, + "empty_like", + lambda *args, **kwargs: original_empty_like(*args, **{**kwargs, "pin_memory": False}), + ) + sampler, _ = cpu_sampler + copier = AsyncWorkerMixin() + copies = DeferredCopies() + copier._async_worker = copies if deferred else None + copier._async_worker_futures = [] + sampler._copy_to_host = copier._copy_to_host + sampler._record_sampler_event = copier._record_sampler_event + storage = torch.empty(8, dtype=torch.int32) + states = [] + for windows in ([3, 4, 1], [4, 2, 1]): + scheduled, outputs, *_ = _sampling_step() + _publish_policy_window_output( + outputs, + torch.tensor(windows, dtype=torch.int32), + num_contexts=2, + num_generations=3, + output_verify_lens=storage, + ) + states.append(sampler.sample_async(scheduled, outputs, [])) + storage.fill_(99) + copies.finish() + for state in states: + sampler.update_requests(state) + assert [r.py_rewind_len for r in states[0].requests] == [0, 1, 2] + assert [r.py_rewind_len for r in states[1].requests] == [0, 2, 0] + assert states[0].host.verify_lens.data_ptr() != states[1].host.verify_lens.data_ptr() + assert copier._async_worker_futures == [] + if deferred: + assert len(states[0].sampler_event.worker_futures) == 4 + assert len(states[1].sampler_event.worker_futures) == 4 + assert states[0].sampler_event.worker_futures is not states[1].sampler_event.worker_futures + + +@pytest.mark.parametrize("native", [False, True]) +def test_marked_execution_never_reads_planning_property(cpu_sampler, native): + class RequestProxy: + def __init__(self, request): + self.request = request + + def __getattr__(self, name): + if name == "py_verify_len": + raise AssertionError("completed execution must not read pre-forward planning") + return getattr(self.request, name) + + sampler, _ = cpu_sampler + scheduled, outputs, *_ = _sampling_step() + scheduled.context_requests_last_chunk = [ + RequestProxy(r) for r in scheduled.context_requests_last_chunk + ] + scheduled.generation_requests = [RequestProxy(r) for r in scheduled.generation_requests] + if native: + outputs["native_uniform_verify"] = True + else: + outputs["verify_lens"] = torch.tensor([1, 1, 3, 4, 1], dtype=torch.int32) + outputs["verify_lens_in_output_order"] = True + sampler.sample_async(scheduled, outputs, []) + + +@pytest.mark.parametrize("explicit_native", [False, True]) +def test_native_uniform_without_optional_ragged_attribute(cpu_sampler, explicit_native): + sampler, _ = cpu_sampler + scheduled, outputs, *_ = _sampling_step() + for request in scheduled.context_requests_last_chunk + scheduled.generation_requests: + del request.py_verify_len + if explicit_native: + outputs["native_uniform_verify"] = True + + state = sampler.sample_async(scheduled, outputs, []) + sampler.update_requests(state) + + assert state.host.verify_lens is None + assert [request.tokens for request in state.requests] == [[4], [8, 9], [12, 13]] + assert [request.py_num_draft_tokens_verified for request in state.requests] == [0, 3, 3] + assert [request.py_rewind_len for request in state.requests] == [3, 2, 2] + + +@pytest.mark.parametrize("ragged_generation", [0, 1]) +def test_legacy_ragged_rejected_with_missing_optional_attributes(cpu_sampler, ragged_generation): + sampler, _ = cpu_sampler + scheduled, outputs, *_ = _sampling_step() + for request in scheduled.context_requests_last_chunk + scheduled.generation_requests: + del request.py_verify_len + scheduled.generation_requests[ragged_generation].py_verify_len = 2 + + with pytest.raises(ValueError, match="ragged producers"): + sampler.sample_async(scheduled, outputs, [])