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..244a26d12d83 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 spec_metadata.is_ragged_verify: + from tensorrt_llm._torch.speculative.ragged_helpers import row_ids_from_lens + + 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 + 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.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/interface.py b/tensorrt_llm/_torch/speculative/interface.py index dc068824b470..8bfd073223c7 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/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). + # 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 .ragged_helpers 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 .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 + 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 .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, + 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/ragged_helpers.py b/tensorrt_llm/_torch/speculative/ragged_helpers.py new file mode 100644 index 000000000000..3c8b6d7826ec --- /dev/null +++ b/tensorrt_llm/_torch/speculative/ragged_helpers.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/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 030fb18f4fd9..5fb5b185dc0b 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) @@ -338,10 +341,30 @@ 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) + 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 @@ -365,16 +388,18 @@ 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 = ( + token_windows[req_idx] - 1 if token_windows is not None else runtime_draft_len ) - 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 +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 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: legacy True is rejected num_context_logits_prefix_sum: Prefix sum of context logits (unused) Returns: @@ -424,6 +453,42 @@ 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") + 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 != 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] + # 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 +531,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() 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_ragged_helpers.py b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_helpers.py new file mode 100644 index 000000000000..65a905d947c0 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_helpers.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.ragged_helpers 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_ragged_metadata.py b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_metadata.py new file mode 100644 index 000000000000..f655f00ed23c --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_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.ragged_helpers 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..7afbea6046a3 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py @@ -0,0 +1,1147 @@ +# 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_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(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]) +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(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): + sampler.sample_async(scheduled, outputs, []) + + +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_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]) +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_verify_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", + [("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) + outputs["verify_lens_in_output_order"] = 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 context.py_rewind_len == (0 if authority == "device" else 3) + 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 + outputs["verify_lens_in_output_order"] = True + 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["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 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 + 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) + 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. + 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) + outputs["verify_lens_in_output_order"] = True + 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) + 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 + assert complete.tokens == [] + assert replacement.tokens == [] + 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, []) 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,